TOPP-RA and joint trajectory animation

airo_drake/__init__.py

@@ -9,7 +9,7 @@
 from airo_drake.building.meshcat import add_meshcat
 from airo_drake.scene import DualArmScene, SingleArmScene
 from airo_drake.visualization.frame import visualize_frame
-from airo_drake.visualization.joint_configurations import animate_joint_configurations
+from airo_drake.visualization.joints import animate_joint_configurations
 
 __all__ = [
     "add_floor",

airo_drake/time_parametrization/toppra.py

@@ -0,0 +1,48 @@
+import numpy as np
+from airo_typing import JointPathType
+from pydrake.multibody.optimization import CalcGridPointsOptions, Toppra
+from pydrake.multibody.plant import MultibodyPlant
+from pydrake.trajectories import PathParameterizedTrajectory, PiecewisePolynomial, Trajectory
+
+
+def time_parametrize_toppra(
+    plant: MultibodyPlant,
+    joints: JointPathType | Trajectory,
+    joint_speed_limit: float = 2.0,  # Max 180 degrees/s ~ 3.14 rad/s
+    joint_acceleration_limit: float = 4.0,  # UR recommends < 800 degrees/s^2 ~ 13.9 rad/s^2
+) -> PathParameterizedTrajectory:
+    """Recalculate the timing of a path or trajectory to respect joint speed and acceleration limits using TOPP-RA.
+
+    Args:
+        plant: The MultibodyPlant for the robot.
+        joints: A joint path or trajectory.
+        joint_speed_limit: The maximum joint speed in rad/s.
+        joint_acceleration_limit: The maximum joint acceleration in rad/s^2.
+    """
+    if isinstance(joints, Trajectory):
+        joint_trajectory = joints
+        n_dofs = joint_trajectory.value(0).shape[0]
+    else:
+        joints = np.array(joints).squeeze()
+        n_dofs = joints.shape[1]
+        times_dummy = np.linspace(0.0, 1.0, len(joints))  # TOPP-RA will calculate the actual times
+        joint_trajectory = PiecewisePolynomial.FirstOrderHold(times_dummy, joints.T)
+
+    gridpoints = Toppra.CalcGridPoints(joint_trajectory, CalcGridPointsOptions())
+
+    acceleration_limits_lower = np.array([-joint_acceleration_limit] * n_dofs)
+    acceleration_limits_upper = np.array([joint_acceleration_limit] * n_dofs)
+    velocity_limits_lower = np.array([-joint_speed_limit] * n_dofs)
+    velocity_limits_upper = np.array([joint_speed_limit] * n_dofs)
+
+    toppra = Toppra(joint_trajectory, plant, gridpoints)
+    toppra.AddJointAccelerationLimit(acceleration_limits_lower, acceleration_limits_upper)
+    toppra.AddJointVelocityLimit(velocity_limits_lower, velocity_limits_upper)
+    time_parametrization = toppra.SolvePathParameterization()
+
+    if time_parametrization is None:
+        raise ValueError("TOPP-RA failed to find a valid time parametrization.")
+
+    joint_trajectory = PathParameterizedTrajectory(joint_trajectory, time_parametrization)
+
+    return joint_trajectory

airo_drake/trajectory/discretize.py

@@ -0,0 +1,14 @@
+import numpy as np
+from airo_typing import JointPathContainer, SingleArmTrajectory
+from pydrake.trajectories import Trajectory
+
+
+def discretize_drake_joint_trajectory(joint_trajectory: Trajectory, steps: int = 100) -> SingleArmTrajectory:
+    positions = []
+    times_uniform = np.linspace(joint_trajectory.start_time(), joint_trajectory.end_time(), steps)
+    for t in times_uniform:
+        position = joint_trajectory.value(t).squeeze()
+        positions.append(position)
+    joint_path = JointPathContainer(positions=np.array(positions))
+    joint_trajectory_discretized = SingleArmTrajectory(times_uniform, joint_path)
+    return joint_trajectory_discretized

airo_drake/visualization/joint_configurations.py → airo_drake/visualization/joints.py

@@ -1,59 +1,98 @@
 import numpy as np
-from airo_typing import JointPathType
+from airo_typing import JointPathType, SingleArmTrajectory
 from loguru import logger
 from pydrake.geometry import Meshcat
 from pydrake.multibody.tree import ModelInstanceIndex
 from pydrake.planning import RobotDiagram
 from pydrake.systems.framework import Context
-from pydrake.trajectories import PiecewisePolynomial
+from pydrake.trajectories import PiecewisePolynomial, Trajectory
 
 
-def animate_joint_configurations(
+def animate_joint_trajectory(
     meshcat: Meshcat,
     robot_diagram: RobotDiagram,
     arm_index: ModelInstanceIndex,
-    joint_configurations: list[np.ndarray] | JointPathType,
-    duration: float = 4.0,
-    time_per_configuration: float = None,
+    joint_trajectory: SingleArmTrajectory | Trajectory,
     context: Context | None = None,
 ) -> None:
-    """Publish a recorded animation to meshcat where the robot arm cycles through the provided joint configurations.
+    """Publish a recorded animation to meshcat where the robot arm follows the provided joint trajectory.
 
     Args:
         meshcat: The MeshCat instance to add the visualization to.
         robot_diagram: The robot diagram.
-        joint_cofigurations: A list of joint configurations to animate.
+        joint_trajectory: The joint trajectory to animate.
         duration: Total duration of the animation.
-        time_per_configuration: The time to spend (still) on each joint solution.  Will overwrite duration if provided.
         context: The context to use for the animation. If None, a new default context will be used (with all other joint angles equal to 0).
     """
-    if joint_configurations is None or len(joint_configurations) == 0:
-        logger.warning("Asked to animate joint configurations, but none provided.")
+    if joint_trajectory is None:
+        logger.warning("Asked to animate joint trajectory, but none provided.")
+
+    if isinstance(joint_trajectory, SingleArmTrajectory):
+        positions = joint_trajectory.path.positions
+        times = joint_trajectory.times
+        if positions is None:
+            logger.warning("Asked to animate SingleArmTrajectory, but it has not positions.")
+            return
+        joint_trajectory = PiecewisePolynomial.FirstOrderHold(times, positions.T)  # linear interpolation
+
+    assert isinstance(joint_trajectory, Trajectory)  # For mypy, but extra check we have a Drake Trajectory now
 
     if context is None:
         # Create a new context just for the recording
         context = robot_diagram.CreateDefaultContext()
     plant = robot_diagram.plant()
     plant_context = plant.GetMyContextFromRoot(context)
 
-    joint_configurations = np.array(joint_configurations).squeeze()
+    fps = 60.0
 
-    if time_per_configuration is not None:
-        duration = time_per_configuration * len(joint_configurations)
+    duration = joint_trajectory.end_time()
 
-    joint_times = np.linspace(0.0, duration, len(joint_configurations))
-    joints_interpolated = PiecewisePolynomial.ZeroOrderHold(joint_times, joint_configurations.T)
+    # TODO can support trajectories that don't start at 0?
+    if not np.isclose(joint_trajectory.start_time(), 0.0):
+        logger.warning("Joint trajectory does not start at time 0.0, this is not officially supported.")
 
-    fps = 60.0
     n_frames = int(max(duration * fps, 1))
     frame_times = np.linspace(0.0, duration, n_frames)
 
     meshcat.StartRecording(set_visualizations_while_recording=False, frames_per_second=fps)
 
     for t in frame_times:
         context.SetTime(t)
-        plant.SetPositions(plant_context, arm_index, joints_interpolated.value(t).squeeze())
+        plant.SetPositions(plant_context, arm_index, joint_trajectory.value(t).squeeze())
         robot_diagram.ForcedPublish(context)
 
     meshcat.StopRecording()
     meshcat.PublishRecording()
+
+
+def animate_joint_configurations(
+    meshcat: Meshcat,
+    robot_diagram: RobotDiagram,
+    arm_index: ModelInstanceIndex,
+    joint_configurations: list[np.ndarray] | JointPathType,
+    duration: float = 4.0,
+    time_per_configuration: float | None = None,
+    context: Context | None = None,
+) -> None:
+    """Publish a recorded animation to meshcat where the robot arm cycles through the provided joint configurations.
+
+    Args:
+        meshcat: The MeshCat instance to add the visualization to.
+        robot_diagram: The robot diagram.
+        joint_cofigurations: A list of joint configurations to animate.
+        duration: Total duration of the animation.
+        time_per_configuration: The time to spend (still) on each joint solution.  Will overwrite duration if provided.
+        context: The context to use for the animation. If None, a new default context will be used (with all other joint angles equal to 0).
+    """
+    if joint_configurations is None or len(joint_configurations) == 0:
+        logger.warning("Asked to animate joint configurations, but none provided.")
+
+    joint_configurations = np.array(joint_configurations).squeeze()
+
+    if time_per_configuration is not None:
+        duration = time_per_configuration * len(joint_configurations)
+
+    joint_times = np.linspace(0.0, duration, len(joint_configurations))
+    joints_interpolated = PiecewisePolynomial.ZeroOrderHold(joint_times, joint_configurations.T)
+
+    animate_joint_trajectory(meshcat, robot_diagram, arm_index, joints_interpolated, context)

notebooks/02_visualization.ipynb

@@ -156,8 +156,7 @@
     "\n",
     "with np.printoptions(precision=3, suppress=True):\n",
     "    for joint_solution in joint_solutions:\n",
-    "        print(joint_solution)\n",
-    "\n"
+    "        print(joint_solution)"
    ]
   },
   {
@@ -217,21 +216,119 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from airo_drake import animate_joint_configurations\n",
-    "\n",
     "animate_joint_configurations(scene.meshcat, scene.robot_diagram, scene.arm_index, joint_path, context=context)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 5. Visualizing trajectories 🌠"
+    "## 5. Visualizing trajectories 🌠\n",
+    "\n",
+    "Path don't have times associated with them (explicitly at least).\n",
+    "Adding times to a path is called *time parameterization*.\n",
+    "Drake provides one that works respects joint velocity and acceleration limits called [TOPP-RA](https://drake.mit.edu/doxygen_cxx/classdrake_1_1multibody_1_1_toppra.html).\n",
+    "\n",
+    "We provide a convenience function for using TOPP-RA with the `airo-mono` path and trajectory types:"
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_drake.time_parametrization.toppra import time_parametrize_toppra\n",
+    "import airo_models\n",
+    "from pydrake.planning import RobotDiagram\n",
+    "\n",
+    "\n",
+    "def create_robot_diagram_without_gripper() -> RobotDiagram:\n",
+    "    robot_diagram_builder = RobotDiagramBuilder()\n",
+    "    parser = robot_diagram_builder.parser()\n",
+    "    plant = robot_diagram_builder.plant()\n",
+    "    arm_urdf_path = airo_models.get_urdf_path(\"ur5e\")\n",
+    "    arm_index = parser.AddModels(arm_urdf_path)[0]\n",
+    "    world_frame = plant.world_frame()\n",
+    "    arm_frame = plant.GetFrameByName(\"base_link\", arm_index)\n",
+    "    plant.WeldFrames(world_frame, arm_frame)\n",
+    "    robot_diagram = robot_diagram_builder.Build()\n",
+    "    return robot_diagram\n",
+    "\n",
+    "\n",
+    "robot_diagram_without_gripper = create_robot_diagram_without_gripper()\n",
+    "\n",
+    "print(\"DoFs in scene without gripper:\", robot_diagram_without_gripper.plant().num_positions())\n",
+    "joint_trajectory = time_parametrize_toppra(\n",
+    "    robot_diagram_without_gripper.plant(), joint_path, joint_acceleration_limit=1.0\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_drake.trajectory.discretize import discretize_drake_joint_trajectory\n",
+    "\n",
+    "\n",
+    "joint_trajectory_discretized = discretize_drake_joint_trajectory(joint_trajectory)\n",
+    "\n",
+    "plt.figure(figsize=(20, 6))\n",
+    "plt.subplot(1, 2, 1)\n",
+    "plt.title(\"Joint angles - before TOPP-RA\")\n",
+    "for row in joint_path.T:\n",
+    "    plt.plot(row)\n",
+    "plt.ylim(-np.pi, np.pi)\n",
+    "plt.subplot(1, 2, 2)\n",
+    "plt.title(\"Joint angles - after TOPP-RA\")\n",
+    "for row in joint_trajectory_discretized.path.positions.T:\n",
+    "    plt.plot(joint_trajectory_discretized.times, row)\n",
+    "plt.ylim(-np.pi, np.pi)\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_drake.visualization.joints import animate_joint_trajectory\n",
+    "\n",
+    "animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, joint_trajectory, context)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Just to show the same function works with our airo-mono discretized trajectory\n",
+    "animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, joint_trajectory_discretized, context)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from airo_typing import JointPathContainer, SingleArmTrajectory\n",
+    "\n",
+    "dance_trajectory = SingleArmTrajectory(\n",
+    "    times=np.linspace(0, 10, 10), path=JointPathContainer(positions=np.random.uniform(-2 * np.pi, 2 * np.pi, (10, 6)))\n",
+    ")\n",
+    "animate_joint_trajectory(scene.meshcat, scene.robot_diagram, scene.arm_index, dance_trajectory, context)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": []
   }
  ],

pyproject.toml

@@ -18,7 +18,7 @@ classifiers = [
   "Operating System :: OS Independent",
 ]
 
-dependencies = ["airo-models", "drake", "ur-analytic-ik", "loguru"]
+dependencies = ["airo-models", "drake", "ur-analytic-ik>=0.0.4", "loguru"]
 
 [project.urls]
 Homepage = "https://github.com/airo-ugent/airo-drake"