diff --git a/airo_drake/__init__.py b/airo_drake/__init__.py index a16f646..074d306 100644 --- a/airo_drake/__init__.py +++ b/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", diff --git a/airo_drake/time_parametrization/toppra.py b/airo_drake/time_parametrization/toppra.py new file mode 100644 index 0000000..241b9d4 --- /dev/null +++ b/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 diff --git a/airo_drake/trajectory/discretize.py b/airo_drake/trajectory/discretize.py new file mode 100644 index 0000000..81263aa --- /dev/null +++ b/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 diff --git a/airo_drake/visualization/joint_configurations.py b/airo_drake/visualization/joints.py similarity index 52% rename from airo_drake/visualization/joint_configurations.py rename to airo_drake/visualization/joints.py index a7a0360..8a74120 100644 --- a/airo_drake/visualization/joint_configurations.py +++ b/airo_drake/visualization/joints.py @@ -1,34 +1,41 @@ 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 @@ -36,15 +43,14 @@ def animate_joint_configurations( 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) @@ -52,8 +58,41 @@ def animate_joint_configurations( 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) diff --git a/notebooks/02_visualization.ipynb b/notebooks/02_visualization.ipynb index 9bdc59a..5c40804 100644 --- a/notebooks/02_visualization.ipynb +++ b/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,8 +216,6 @@ "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)" ] }, @@ -226,12 +223,112 @@ "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": [] } ], diff --git a/pyproject.toml b/pyproject.toml index c641a2f..e400e69 100644 --- a/pyproject.toml +++ b/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"