Source code for brom_drake.example_helpers.block_handler_system

from importlib import resources as impresources

from pydrake.math import RollPitchYaw, RigidTransform, RotationMatrix
from pydrake.multibody.math import SpatialVelocity
from pydrake.multibody.parsing import Parser
from pydrake.multibody.plant import MultibodyPlant, CoulombFriction
from pydrake.systems.framework import Diagram, DiagramBuilder, LeafSystem, PortDataType, BasicVector

import numpy as np

# Internal Imports
from brom_drake import example_helpers as eh
from brom_drake.utils import AddGround, AddMultibodyTriad



class BlockHandlerSystem(LeafSystem):
    def __init__(
        self,
        plant: MultibodyPlant,
        scene_graph,
        block_name: str = "block_with_slots",
        model_urdf_path: str = None,
    ):
        LeafSystem.__init__(self)

        # Input Processing
        if model_urdf_path is None:
            model_urdf_path = str(
                impresources.files(eh) / "models/slider-block.urdf",
            )

        # Constants
        self.block_name = block_name

        # Add the Block to the given plant
        self.plant = plant
        self.block_model_idx = Parser(plant=self.plant).AddModels(model_urdf_path)[0]
        self.block_model_name = self.plant.GetModelInstanceName(self.block_model_idx)
        self.block_body_name = "block"

        AddGround(self.plant) #Add ground to plant

        # Add the Triad
        self.scene_graph = scene_graph
        AddMultibodyTriad(plant.GetFrameByName(self.block_body_name), self.scene_graph)

        self.plant.Finalize()
        self.context = self.plant.CreateDefaultContext()

        # Create Input Port for the Slider Block System
        self.desired_pose_port = self.DeclareVectorInputPort(
            "desired_pose",
            BasicVector(6),
        )

        # Create Output Port which should share the pose of the block
        self.DeclareVectorOutputPort(
            "measured_block_pose",
            BasicVector(6),
            self.SetBlockPose,
            {self.time_ticket()}    # indicate that this doesn't depend on any inputs,
        )                           # but should still be updated each timestep



    def SetBlockPose(self, context, output):
        """
        Description:
            This function sets the desired pose of the block.
        """

        # Get Desired Pose from Port
        plant_context = self.context
        pose_as_vec = self.desired_pose_port.Eval(context)

        self.plant.SetFreeBodyPose(
            plant_context,
            self.plant.GetBodyByName(self.block_body_name),
            RigidTransform(RollPitchYaw(pose_as_vec[:3]),pose_as_vec[3:])
        )

        # self.plant.SetFreeBodySpatialVelocity(
        #     self.plant.GetBodyByName(self.block_body_name),
        #     SpatialVelocity(np.zeros(3),np.array([0.0,0.0,0.0])),
        #     plant_context
        # )

        X_WBlock = self.plant.GetFreeBodyPose(
            plant_context,
            self.plant.GetBodyByName(self.block_body_name)
        )

        pose_as_vector = np.hstack([RollPitchYaw(X_WBlock.rotation()).vector(), X_WBlock.translation()])

        # Create Output
        output.SetFromVector(pose_as_vector)




    def SetInitialBlockState(self,diagram_context):
        """
        Description:
            Sets the initial position to be slightly above the ground (small, positive z value)
            to be .
        """

        # Set Pose
        p_WBlock = [0.0, 0.0, 0.2]
        R_WBlock = RotationMatrix.MakeXRotation(np.pi/2.0) # RotationMatrix.MakeXRotation(-np.pi/2.0)
        X_WBlock = RigidTransform(R_WBlock, p_WBlock)
        self.plant.SetFreeBodyPose(
            self.plant.GetMyContextFromRoot(diagram_context),
            self.plant.GetBodyByName(self.block_body_name),
            X_WBlock)

        # Set Velocities
        self.plant.SetFreeBodySpatialVelocity(
            self.plant.GetBodyByName(self.block_body_name),
            SpatialVelocity(np.zeros(3),np.array([0.0,0.0,0.0])),
            self.plant.GetMyContextFromRoot(diagram_context))