Add mobile platform support (#7)

* Add mobile platform support and example notebook

* arm_index is now optional in MobilePlatformWithSingleArmScene

Because it's possible to want to do motion planning without the arm.
Another solution would be to make the arm static, which would also check
for collisions based on the arm itself.

* Implement functionality required for MOBI integration in airo-planner

* Update toppra.py (acc. limits)

* Update mobile_platform.py

Avoid collisions with parts of MOBI

* Update scene.py

* Update scene.py

* Run pre-commit

* Export animate_mobile_platform_trajectory

airo_drake/__init__.py

@@ -11,6 +11,7 @@
 from airo_drake.building.finish import finish_build
 from airo_drake.building.floor import add_floor
 from airo_drake.building.manipulator import X_URBASE_ROSBASE, X_URTOOL0_ROBOTIQ, add_manipulator
+from airo_drake.building.mobile_platform import add_mobile_platform, attach_mobile_platform_to_world
 from airo_drake.building.meshcat import add_meshcat
 from airo_drake.building.wall import add_wall
 from airo_drake.path.analysis import (
@@ -26,10 +27,10 @@
     calculate_valid_joint_paths,
     create_paths_from_closest_solutions,
 )
-from airo_drake.scene import DualArmScene, SingleArmScene
+from airo_drake.scene import DualArmScene, SingleArmScene, MobilePlatformWithSingleArmScene
 from airo_drake.trajectory.timing import shift_drake_trajectory_in_time
 from airo_drake.trajectory.concatenate import concatenate_drake_trajectories
-from airo_drake.time_parametrization.toppra import time_parametrize_toppra
+from airo_drake.time_parametrization.toppra import time_parametrize_toppra, time_parametrize_toppra_mobile_platform
 from airo_drake.trajectory.discretize import discretize_drake_joint_trajectory, discretize_drake_pose_trajectory
 from airo_drake.trajectory.interpolate import joint_trajectory_to_drake
 from airo_drake.visualization.frame import visualize_frame
@@ -38,27 +39,32 @@
     animate_joint_configurations,
     animate_joint_trajectory,
 )
+from airo_drake.visualization.mobile_platform import animate_mobile_platform_trajectory
 from airo_drake.collision.collision_checking import (
     filter_collisions_between_all_body_pairs,
     list_collisions_between_bodies,
 )
 
-
 __all__ = [
     "add_floor",
     "add_wall",
     "add_manipulator",
+    "add_mobile_platform",
+    "attach_mobile_platform_to_world",
     "X_URBASE_ROSBASE",
     "X_URTOOL0_ROBOTIQ",
     "add_meshcat",
     "finish_build",
     "SingleArmScene",
     "DualArmScene",
+    "MobilePlatformWithSingleArmScene",
     "visualize_frame",
     "animate_joint_configurations",
     "animate_joint_trajectory",
     "animate_dual_joint_trajectory",
+    "animate_mobile_platform_trajectory",
     "time_parametrize_toppra",
+    "time_parametrize_toppra_mobile_platform",
     "discretize_drake_joint_trajectory",
     "discretize_drake_pose_trajectory",
     "joint_trajectory_to_drake",

airo_drake/building/manipulator.py

@@ -2,7 +2,7 @@
 import numpy as np
 from airo_typing import HomogeneousMatrixType
 from pydrake.math import RigidTransform, RollPitchYaw
-from pydrake.multibody.tree import ModelInstanceIndex
+from pydrake.multibody.tree import Frame, ModelInstanceIndex
 from pydrake.planning import RobotDiagramBuilder
 
 # X_URBASE_ROSBASE is the 180 rotation between ROS URDF base and the UR control box base
@@ -18,7 +18,9 @@ def add_manipulator(
     gripper_name: str,
     arm_transform: HomogeneousMatrixType | None = None,
     gripper_transform: HomogeneousMatrixType | None = None,
+    static_arm: bool = False,
     static_gripper: bool = False,
+    parent_frame: Frame | None = None,
 ) -> tuple[ModelInstanceIndex, ModelInstanceIndex]:
     """Add a manipulator (a robot arm with a gripper) to the robot diagram builder.
     Looks up the URDF files for the robot and gripper and welds them together.
@@ -32,7 +34,9 @@ def add_manipulator(
         gripper_name: The name of the gripper, must be known by airo-models
         arm_transform: The transform of the robot arm, if None, we use supply a robot-specific default.
         gripper_transform: The transform of the gripper, if None, we supply a default for the robot-gripper pair.
+        static_arm: If True, will fix all arm joints to their default. Useful when you don't want the arm DoFs in the plant.
         static_gripper: If True, will fix all gripper joints to their default. Useful when you don't want the gripper DoFs in the plant.
+        parent_frame: The parent frame to weld to. If None, we use the world frame of the plant.
 
     Returns:
         The robot and gripper index.
@@ -46,6 +50,11 @@ def add_manipulator(
     arm_urdf_path = airo_models.get_urdf_path(name)
     gripper_urdf_path = airo_models.get_urdf_path(gripper_name)
 
+    if static_arm:
+        arm_urdf = airo_models.urdf.read_urdf(arm_urdf_path)
+        airo_models.urdf.make_static(arm_urdf)
+        arm_urdf_path = airo_models.urdf.write_urdf_to_tempfile(arm_urdf, arm_urdf_path, prefix=f"{name}_static_")
+
     arm_index = parser.AddModels(arm_urdf_path)[0]
 
     if static_gripper:
@@ -79,7 +88,9 @@ def add_manipulator(
     else:
         gripper_rigid_transform = RigidTransform(gripper_transform)
 
-    plant.WeldFrames(world_frame, arm_frame, arm_rigid_transform)
+    if parent_frame is None:
+        parent_frame = world_frame
+    plant.WeldFrames(parent_frame, arm_frame, arm_rigid_transform)
     plant.WeldFrames(arm_tool_frame, gripper_frame, gripper_rigid_transform)
 
     return arm_index, gripper_index

airo_drake/building/mobile_platform.py

@@ -0,0 +1,146 @@
+from typing import List, Tuple
+
+import airo_models
+import numpy as np
+from airo_typing import Vector2DType
+from pydrake.math import RigidTransform, RollPitchYaw
+from pydrake.multibody.tree import ModelInstanceIndex, PlanarJoint
+from pydrake.planning import RobotDiagramBuilder
+
+# 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
+
+# 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_mobile_platform(
+    robot_diagram_builder: RobotDiagramBuilder,
+    drive_positions: Tuple[Vector2DType] = (
+        np.array([1, -0.5]),
+        np.array([1, 0.5]),
+        np.array([-1, -0.5]),
+        np.array([-1, 0.5]),
+    ),
+    battery_position: Vector2DType = np.array([0, 0.5]),
+    cpu_position: Vector2DType = np.array([0, -0.5]),
+    side_height: float = 0.43,
+    side_length: float = 0.69,
+    roof_width: float = 0.525,
+    roof_thickness: float = 0.03,
+    brick_size: float = 0.233,
+) -> Tuple[ModelInstanceIndex, List[ModelInstanceIndex]]:
+    """Add a mobile platform on wheels to the robot diagram builder.
+
+    Currently, this method specifically loads the KELO Robile platform bricks.
+    It has default values that correspond to the configuration of the platform at AIRO.
+    In the future, this may change.
+
+    Looks up the URDF files for the bricks and welds them together. Also adds the robot's frame, approximated with
+    boxes. This method returns the model instance index of a frame, which can be used to weld the robot to the
+    world frame, or to add a PlanarJoint to make it possible to move the robot around.
+
+    Args:
+        robot_diagram_builder: The robot diagram builder to which the platform will be added.
+        drive_positions: A tuple of 2D positions of the drive bricks, relative to the brick size and root frame.
+        battery_position: A 2D position of the battery brick, relative to the brick size and root frame.
+        cpu_position: A 2D position of the CPU brick, relative to the brick size and root frame.
+        side_height: The height of the mobile robot's side.
+        side_length: The length (along X) of the mobile robot.
+        roof_width: The width (along Y) of the mobile robot.
+        roof_thickness: The thickness of the roof.
+        brick_size: The size (width and length) of a KELO brick.
+
+    Returns:
+        The mobile platform index and the indices of the parts that make up the mobile platform.
+    """
+
+    plant = robot_diagram_builder.plant()
+    parser = robot_diagram_builder.parser()
+    parser.SetAutoRenaming(True)
+
+    mobi_urdf_path = airo_models.get_urdf_path("kelo_robile")
+
+    drive_transforms = [
+        RigidTransform(p=[brick_size * p[0], brick_size * p[1], 0], rpy=RollPitchYaw([0, 0, 0]))
+        for p in drive_positions
+    ]
+
+    mobi_model_index = parser.AddModels(mobi_urdf_path)[0]
+    robot_root_frame = plant.GetFrameByName("base_link", mobi_model_index)
+
+    mobi_part_indices = []
+
+    # robot_transform: relative to world
+    # drive_transforms: relative to robot_transform
+    for drive_index, drive_transform in enumerate(drive_transforms):
+        brick_index = parser.AddModels(airo_models.get_urdf_path("kelo_robile_wheel"))[0]
+        brick_frame = plant.GetFrameByName("base_link", brick_index)
+        plant.WeldFrames(robot_root_frame, brick_frame, drive_transform)
+        mobi_part_indices.append(brick_index)
+
+    battery_transform = RigidTransform(
+        p=[brick_size * battery_position[0], brick_size * battery_position[1], 0], rpy=RollPitchYaw([0, 0, 0])
+    )
+    battery_index = parser.AddModels(airo_models.get_urdf_path("kelo_robile_battery"))[0]
+    battery_frame = plant.GetFrameByName("base_link", battery_index)
+    plant.WeldFrames(robot_root_frame, battery_frame, battery_transform)
+    mobi_part_indices.append(battery_index)
+
+    cpu_transform = RigidTransform(
+        p=[brick_size * cpu_position[0], brick_size * cpu_position[1], 0], rpy=RollPitchYaw([0, 0, 0])
+    )
+    cpu_index = parser.AddModels(airo_models.get_urdf_path("kelo_robile_cpu"))[0]
+    cpu_frame = plant.GetFrameByName("base_link", cpu_index)
+    plant.WeldFrames(robot_root_frame, cpu_frame, cpu_transform)
+    mobi_part_indices.append(cpu_index)
+
+    front_brick_position = np.max([brick_size * p[0] for p in drive_positions]) + brick_size / 2
+
+    side_height_half = 0.5 * side_height
+    side_length_half = 0.5 * side_length
+    side_transforms = [
+        RigidTransform(p=[front_brick_position - side_length_half, -brick_size, side_height_half]),
+        RigidTransform(p=[front_brick_position - side_length_half, brick_size, side_height_half]),
+    ]
+
+    for side_transform in side_transforms:
+        side_urdf_path = airo_models.box_urdf_path([side_length, 0.03, side_height], "side")
+        side_index = parser.AddModels(side_urdf_path)[0]
+        side_frame = plant.GetFrameByName("base_link", side_index)
+        plant.WeldFrames(robot_root_frame, side_frame, side_transform)
+        mobi_part_indices.append(side_index)
+
+    roof_length = side_length
+    roof_length_half = 0.5 * roof_length
+    roof_thickness_half = 0.5 * roof_thickness
+    roof_transform = RigidTransform(p=[front_brick_position - roof_length_half, 0, side_height + roof_thickness_half])
+
+    roof_urdf_path = airo_models.box_urdf_path([roof_length, roof_width, 0.03], "roof")
+    roof_index = parser.AddModels(roof_urdf_path)[0]
+    roof_frame = plant.GetFrameByName("base_link", roof_index)
+    plant.WeldFrames(robot_root_frame, roof_frame, roof_transform)
+    mobi_part_indices.append(roof_index)
+
+    return mobi_model_index, mobi_part_indices
+
+
+def attach_mobile_platform_to_world(
+    robot_diagram_builder: RobotDiagramBuilder, mobi_model_index: ModelInstanceIndex, static_platform: bool
+):
+    """Attach the mobile platform to the world frame with a planar joint. This allows the mobile platform to
+    move around by setting the (x, y, theta) values of the joint.
+
+    Args:
+        robot_diagram_builder: The robot diagram builder to which the platform will be added.
+        mobi_model_index: The mobile platform index.
+        static_platform: If true, will weld. If False, will make a joint."""
+    plant = robot_diagram_builder.plant()
+
+    mobi_frame = plant.GetFrameByName("base_link", mobi_model_index)
+
+    if not static_platform:
+        platform_planar_joint = PlanarJoint("mobi_joint", plant.world_frame(), mobi_frame)
+        plant.AddJoint(platform_planar_joint)
+    else:
+        plant.WeldFrames(plant.world_frame(), mobi_frame)

airo_drake/scene.py

@@ -1,10 +1,23 @@
 from dataclasses import dataclass
+from typing import List
 
 from pydrake.geometry import Meshcat
 from pydrake.multibody.tree import ModelInstanceIndex
 from pydrake.planning import RobotDiagram
 
 
+@dataclass
+class MobilePlatformWithSingleArmScene:
+    """The most important objects when using Drake with a single robot arm mounted on a mobile platform."""
+
+    robot_diagram: RobotDiagram
+    mobile_platform_index: ModelInstanceIndex
+    mobile_platform_part_indices: List[ModelInstanceIndex]
+    arm_index: ModelInstanceIndex | None = None
+    gripper_index: ModelInstanceIndex | None = None
+    meshcat: Meshcat | None = None
+
+
 @dataclass
 class SingleArmScene:
     """The most important objects when using Drake with a single robot arm."""

airo_drake/time_parametrization/toppra.py

@@ -5,6 +5,56 @@
 from pydrake.trajectories import PathParameterizedTrajectory, PiecewisePolynomial, Trajectory
 
 
+def time_parametrize_toppra_mobile_platform(
+    plant: MultibodyPlant,
+    poses: JointPathType,  # TODO type
+    linear_velocity_limit: float = 1.0,
+    angular_velocity_limit: float = np.pi / 8,
+    linear_acceleration_limit: float = 0.25,
+    angular_acceleration_limit: float = 0.025,
+) -> PathParameterizedTrajectory:
+    """Recalculate the timing of a path or trajectory to respect linear and angular velocity and acceleration limits
+    using TOPP-RA.
+
+    Args:
+        plant: The MultibodyPlant for the robot.
+        poses: A pose path or trajectory.
+        linear_velocity_limit: The limit on the linear velocity in m/s.
+        angular_velocity_limit: The limit on the angular velocity of the platform in rad/s.
+        linear_acceleration_limit: The limit on the linear acceleration in m/s^2.
+        angular_acceleration_limit: The limit on the angular acceleration in rad/s^2.
+
+    Returns:
+        A time parameterized trajectory."""
+    if isinstance(poses, Trajectory):
+        pose_trajectory = poses
+    else:
+        poses = np.array(poses).squeeze()
+        times_dummy = np.linspace(0.0, 1.0, len(poses))  # TOPP-RA will calculate the actual times
+        pose_trajectory = PiecewisePolynomial.FirstOrderHold(times_dummy, poses.T)
+
+    gridpoints = Toppra.CalcGridPoints(pose_trajectory, CalcGridPointsOptions())
+
+    acceleration_limits_lower = np.array(
+        [-linear_acceleration_limit, -linear_acceleration_limit, -angular_acceleration_limit]
+    )
+    acceleration_limits_upper = -acceleration_limits_lower
+    velocity_limits_lower = np.array([-linear_velocity_limit, -linear_velocity_limit, -angular_velocity_limit])
+    velocity_limits_upper = -velocity_limits_lower
+
+    toppra = Toppra(pose_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.")
+
+    pose_trajectory = PathParameterizedTrajectory(pose_trajectory, time_parametrization)
+
+    return pose_trajectory
+
+
 def time_parametrize_toppra(
     plant: MultibodyPlant,
     joints: JointPathType | Trajectory,

airo_drake/visualization/mobile_platform.py

@@ -0,0 +1,56 @@
+import numpy as np
+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 Trajectory
+
+
+def animate_mobile_platform_trajectory(
+    meshcat: Meshcat,
+    robot_diagram: RobotDiagram,
+    mobile_platform_index: ModelInstanceIndex,
+    pose_trajectory: Trajectory,
+    context: Context | None = None,
+) -> None:
+    """Publish a recorded animation to meshcat where the mobile platform follows the provided pose trajectory.
+
+    Args:
+        meshcat: The MeshCat instance to add the visualization to.
+        robot_diagram: The robot diagram.
+        mobile_platform_index: The index of the platform to animate.
+        pose_trajectory: The joint trajectory to animate.
+        context: The context to use for the animation. If None, a new default context will be used.
+    """
+    if pose_trajectory is None:
+        logger.warning("Asked to animate pose trajectory, but none provided.")
+
+    assert isinstance(pose_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)
+
+    fps = 60.0
+
+    duration = pose_trajectory.end_time()
+
+    # TODO can support trajectories that don't start at 0?
+    if not np.isclose(pose_trajectory.start_time(), 0.0):
+        logger.warning("Pose trajectory does not start at time 0.0, this is not officially supported.")
+
+    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, mobile_platform_index, pose_trajectory.value(t).squeeze())
+        robot_diagram.ForcedPublish(context)
+
+    meshcat.StopRecording()
+    meshcat.PublishRecording()