Pydrake: Creating a Trajectory Source for RigidTransformations

I've been trying to create a TrajectorySource for RigidTransforms to pass into a DifferentialInverseKinematicsIntegrator which only takes in RigidTransforms in its input port.

def createTraj(time, pose):
    times = []
    poses = []
    for step in time:
        times.append(time[step])
        poses.append(pose[step])
        
    return PiecewisePose.MakeLinear(times, poses)

Initially, I tried to directly pass in the output from createTraj above into TrajectorySource but ran into the issue of my trajectory having more than one columns: Failure at systems/primitives/trajectory_source.cc:21 in TrajectorySource(): condition 'trajectory.cols() == 1' failed.

import matplotlib.pyplot as plt, mpld3
class DexterTest():
# Output from createTraj is passed as parameter: traj into constructor
    def __init__(self, traj):
        builder = DiagramBuilder()
        self.station = DexterPPStation(1e-4, "/opt/drake/share/drake/manipulation/models/final_dexter_description/urdf/dexter.urdf")

        self.station.CreateBins("/opt/drake/share/drake/examples/manipulation_station/models/bin.sdf", RigidTransform(np.array([0.5,0,0])), RigidTransform(np.array([0,0.5,0])))
        self.station.CreateRandomPickingObjects(3)
        self.station.AddDexter()
        builder.AddSystem(self.station)
        self.station.Finalize()
        
        self.diff_ik = DifferentialInverseKinematicsIntegrator(self.station.controller_plant, self.station.plant.GetFrameByName("link6", self.station.dexter["instance"]), self.station.time_step, DifferentialInverseKinematicsParameters(7,7))
        builder.AddSystem(self.diff_ik)
#=========================================== Likely Source of Error ===========================================
        pose = builder.AddSystem(PoseSystem())
        p_G_source = builder.AddSystem(TrajectorySource(traj.get_position_trajectory()))
        w_G_source = builder.AddSystem(TrajectorySource(traj.get_orientation_trajectory()))

        builder.Connect(p_G_source.get_output_port(), pose.GetInputPort("p_G"))
        builder.Connect(w_G_source.get_output_port(), pose.GetInputPort("r_G"))

        builder.Connect(pose.get_output_port(), self.diff_ik.get_input_port())
#======================================================================================           
        MeshcatVisualizerCpp.AddToBuilder(builder, self.station.GetOutputPort("query_object"), meshcat)
        self.diagram = builder.Build()
        self.simulator = Simulator(self.diagram)
        self.diagram_context = self.simulator.get_mutable_context()
        self.station_context = self.station.GetMyMutableContextFromRoot(self.diagram_context)
        self.plant_context = self.station.GetSubsystemContext(self.station.plant, self.station_context)
        self.station.SetRandomPoses(self.plant_context)
        builder.Connect(self.diff_ik.get_output_port(), self.station.GetInputPort("dexter_position"))
        
    def run(self):
        self.simulator.set_target_realtime_rate(2.0)
        self.simulator.AdvanceTo(1)
        
class PoseSystem(LeafSystem):
    def __init__(self):
        LeafSystem.__init__(self)
        
        self.p_G = self.DeclareVectorInputPort("p_G", BasicVector(3))
        self.r_G = self.DeclareVectorInputPort("r_G", BasicVector(4))
        self.DeclareAbstractOutputPort("X_G", Value[RigidTransform], self.CalcOutput)
        
        
    def CalcOutput(self, context, output):
        pose = RigidTransform(Quaternion(self.r_G.Eval(context)), self.p_G.Eval(context))
        output.set_value(pose)
                                          

Instead, I tried to break up my trajectory into its orientation and position parts, add them to the input ports of a custom system, and then reconstruct them together in the output port. However, this gives me the following RuntimeError once the run method is called: RuntimeError: This multibody element does not belong to the supplied MultibodyTree.

Any help would be greatly appreciated!

Solution 1:

I think you are very close. The PoseSystem looks like it should be a solution to the problem you've articulated in your post. (The error about MultibodyTree must be coming from the other part of your code.

You don't actually need to break the RigidTransform up into orientation / translation to create your PoseSystem, your CalcOutput could just call output.set_value(poses.Eval(t)) if poses is a PiecewisePose trajectory.

I have an example of doing this in the PickAndPlaceTrajectory class in this notebook: https://github.com/RussTedrake/manipulation/blob/008cec6343dd39063705287e6664a3fee71a43b8/pose.ipynb