I have a part object that I've loaded into the drake simulator. The part is a free body, with a floating joint to the end effector. The end effector is a suction joint, which we can lock/unlock.

I want to initialize the part on the table, so that we can pick it up from the table later. I want to specify the part's start position in terms of the world frame.

However, when I use plant.SetFreeBodyPose(plant_sim_context, part_body, world_to_part), the part seems to be transformed relative to the end effector's frame, not the world frame.

I have to compute the relative transform from the end effector to the part, and use plant.SetFreeBodyPose(plant_sim_context, part_body, ee_to_part) so that the part is actual transformed relative to the world frame.

Based on the documentation, I thought SetFreeBodyPose is supposed to be relative to the world frame:

Sets context to store the pose X_WB of a given body B in the world frame W.

So I'm confused why I need to calculate the relative transform to the end effector.

Here is my full code for reference:

from pathlib import Path
from time import time, sleep

from pydrake.all import (
    CollisionCheckerContext,
    Meshcat,
    RigidTransform,
    RobotDiagramBuilder,
    RollPitchYaw,
    SceneGraphCollisionChecker,
    QuaternionFloatingJoint,
    Quaternion,
)
from pydrake.systems.analysis import Simulator
from pydrake.visualization import AddDefaultVisualization

from pivot_autogrind.config import COMBO_END_EFFECTOR_URDF_PATH
import numpy as np

from pivot_autogrind.visualization.simulator_visualizer import SimulatorVisualizer

ROOT_PATH = Path(__file__).parent.parent.parent.parent.as_posix()
ROBOT_URDF_PATH = ROOT_PATH + "/component_configs/abb/urdf/irb1200_7_70.urdf"
PART1_URDF_PATH = ROOT_PATH + "/resources/test_objects/part1.urdf"
TABLE_URDF_PATH = ROOT_PATH + "/component_configs/ingress_table/urdf/ingress_table.urdf"


def add_suction_joint(plant):
    suction_body = plant.GetBodyByName("combo_end_effector_link")

    object_id = plant.GetModelInstanceByName("part1")
    object_body = plant.GetBodyByName("part1_link", object_id)
    suction_joint = plant.AddJoint(
        QuaternionFloatingJoint(
            "suction_joint",
            suction_body.body_frame(),
            object_body.body_frame(),
        ),
    )

    # Position the suction joint to align with the long axis of the object
    translation = [0.03, 0, 0.34]
    rotation = RollPitchYaw(
        0,
        np.pi / 2,
        0,
    ).ToQuaternion() 
    suction_to_object = RigidTransform(rotation, translation)

    suction_joint.SetDefaultPose(suction_to_object)
    return suction_joint


def launch_abb_simulator(
    time_step: float,
    realtime_rate: float,
    set_collision_filter: bool = False,
) -> None:
    robot_diagram_builder = RobotDiagramBuilder(time_step=0)
    builder = robot_diagram_builder.builder()
    plant = robot_diagram_builder.plant()
    parser = robot_diagram_builder.parser()

    # Add robot
    robot_model = parser.AddModels(ROBOT_URDF_PATH)
    plant.WeldFrames(plant.world_frame(), plant.GetFrameByName("base_link"))

    # Add the end effector
    end_effector_model = parser.AddModels(COMBO_END_EFFECTOR_URDF_PATH)
    translation = [0.01, 0, 0]
    rotation = RollPitchYaw(
        np.pi / 2,
        0,
        np.pi / 2,
    ).ToQuaternion()  # This rotation makes the gripper oriented along the axis of the final joint
    ee_transform = RigidTransform(rotation, translation)
    plant.WeldFrames(plant.GetFrameByName("link_6"), plant.GetFrameByName("combo_end_effector_link"), ee_transform)


    # Add the part object
    part_model = parser.AddModels(PART1_URDF_PATH)
    suction_joint = add_suction_joint(plant)

    # Finalize the scene and add visualization
    plant.Finalize()
    meshcat = Meshcat()
    AddDefaultVisualization(builder, meshcat)
    diagram = robot_diagram_builder.Build()
    simulator = Simulator(diagram)

    # Set context
    simulator_context = simulator.get_mutable_context()
    plant_sim_context = plant.GetMyContextFromRoot(simulator_context)

    # Unlock the suction joint
    suction_joint.Unlock(plant_sim_context)

    # Set the initial positions
    robot_joint_positions = [0, 0, 0, 0, 0, 0]
    part_quaternion = [1, 0, 0, 0]
    part_translation = [0, 0, 0]
    positions = robot_joint_positions + part_quaternion + part_translation
    plant.SetPositions(plant_sim_context, positions)

    
    part_body = plant.GetBodyByName("part1_link")

    # World to end effector transform
    world_to_ee = plant.CalcRelativeTransform(
        plant_sim_context,
        plant.world_frame(),
        plant.GetFrameByName("combo_end_effector_link"),
    )

    # Part pose in world frame
    world_to_part = RigidTransform(pose=[0, 0.5, 0])

    # Part pose in end effector frame 
    ee_to_part = world_to_ee.InvertAndCompose(world_to_part)

    # Set the body pose
    plant.SetFreeBodyPose(plant_sim_context, part_body, ee_to_part)

    simulator.Initialize()
    simulator.set_publish_every_time_step(False)
    simulator.set_target_realtime_rate(realtime_rate)

    sleep(60)

if __name__ == "__main__":
    launch_abb_simulator(
        time_step=3e-3,
        realtime_rate=1.0,
        set_collision_filter=True,
    )