Add mobile platform support and example notebook

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,7 +27,7 @@
     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
@@ -48,12 +49,15 @@
     "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",

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 ModelInstanceIndex, Frame
 from pydrake.planning import RobotDiagramBuilder
 
 # X_URBASE_ROSBASE is the 180 rotation between ROS URDF base and the UR control box base
@@ -13,12 +13,13 @@
 
 
 def add_manipulator(
-    robot_diagram_builder: RobotDiagramBuilder,
-    name: str,
-    gripper_name: str,
-    arm_transform: HomogeneousMatrixType | None = None,
-    gripper_transform: HomogeneousMatrixType | None = None,
-    static_gripper: bool = False,
+        robot_diagram_builder: RobotDiagramBuilder,
+        name: str,
+        gripper_name: str,
+        arm_transform: HomogeneousMatrixType | None = None,
+        gripper_transform: HomogeneousMatrixType | None = None,
+        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.
@@ -33,6 +34,7 @@ def add_manipulator(
         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_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.
@@ -79,7 +81,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,127 @@
+from typing import 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
+) -> 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.
+    """
+
+    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)
+
+    # 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)
+
+    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)
+
+    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)
+
+    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)
+
+    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)
+
+    return mobi_model_index
+
+
+def attach_mobile_platform_to_world(robot_diagram_builder: RobotDiagramBuilder, mobi_model_index: ModelInstanceIndex):
+    """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."""
+    plant = robot_diagram_builder.plant()
+
+    mobi_frame = plant.GetFrameByName("base_link", mobi_model_index)
+
+    platform_planar_joint = PlanarJoint("mobi_joint", plant.world_frame(), mobi_frame)
+    plant.AddJoint(platform_planar_joint)

airo_drake/scene.py

@@ -5,6 +5,17 @@
 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
+    arm_index: ModelInstanceIndex
+    gripper_index: ModelInstanceIndex | None = None
+    meshcat: Meshcat | None = None
+
+
 @dataclass
 class SingleArmScene:
     """The most important objects when using Drake with a single robot arm."""