from importlib import resources as impresources
import numpy as np
from pydrake.all import (
AddMultibodyPlantSceneGraph,
ConstantVectorSource,
Frame,
InverseDynamicsController,
ModelInstanceIndex,
MultibodyPlant,
Multiplexer,
Parser,
RigidTransform,
)
from pydrake.geometry import SceneGraph, Meshcat, MeshcatVisualizer
from pydrake.math import RollPitchYaw, RigidTransform, RotationMatrix
from pydrake.multibody.tree import FixedOffsetFrame
from pydrake.systems.framework import Diagram, DiagramBuilder
from pydrake.systems.primitives import Demultiplexer
from typing import Tuple
# Local imports
from brom_drake.robots.gripper_type import GripperType
from brom_drake.control.grippers.gripper_controller import GripperController
from brom_drake.file_manipulation.urdf import DrakeReadyURDFConverter
from brom_drake.utils.leaf_systems import EndEffectorWrenchCalculator
from brom_drake import robots
[docs]
class UR10eStation(Diagram):
"""
*Description*
A template system diagram for controlling a UR10e robot in a simulated environment.
*Diagram*
The diagram's inputs and outputs can be represented as: ::
-----------------
| |
gripper_target -------------------->| |
(Vector, np.array) | UR10e |
| |
gripper_target_type --------------->| Station |
(Abstract, GripperTarget Enum) | |
| |
desired_joint_positions ----------->| |
(Vector, np.array) | |
| |
-----------------
"""
def __init__(
self,
time_step: float = 0.002,
meshcat_port_number: int = 7001,
gripper_type: GripperType = GripperType.Robotiq_2f_85,
):
"""
*Description*
This class defines the UR10e Station.
"""
# Input Processing
self.meshcat_port_number = meshcat_port_number
self.time_step = time_step
# Initialize the Diagram
Diagram.__init__(self)
self.set_name("UR10e_Station")
# Create a diagram builder
self.builder = DiagramBuilder()
# Create the plant and scene graph
self.plant, self.scene_graph = None, None
self.controller_plant = None
self.create_plants_and_scene_graph(time_step)
# Body ID's and Poses for Anything else in the scene
self.object_ids = []
self.object_poses = []
# Whether we have a camera in the simulation
# self.has_camera = False
self.AddArm()
# Which sort of gripper we're using (if any)
self.gripper_type = gripper_type
self.gripper_controller = None
if gripper_type == GripperType.Robotiq_2f_85:
self.Add2f85Gripper()
self.add_gripper_controller()
[docs]
def add_arm_to_plant_with_ee_frame(
self,
arm_urdf_path: str,
plant: MultibodyPlant,
new_frame_name: str = "end_effector",
X_ee: RigidTransform = RigidTransform(),
) -> Tuple[ModelInstanceIndex, Frame]:
"""
*Description*
This function adds the "arm" object at the arm_urdf_path location to the plant.
*Parameters*
arm_urdf_path : str
The path to the URDF file for the arm.
plant : MultibodyPlant
The plant to which the arm will be added.
*Returns*
arm: ModelInstanceIndex
The model instance index of the arm and the frame of the end effector.
end_effector_frame: Frame
The frame of the end effector.
"""
# Setup
# Add the arm to the provided plant
arm = Parser(plant=plant).AddModels(arm_urdf_path)[0]
# Fix the base of the arm to the world
plant.WeldFrames(
plant.world_frame(),
plant.GetFrameByName("base_link", arm),
)
# Add Frame for the end effector
end_effector_frame = plant.AddFrame(
FixedOffsetFrame(
new_frame_name,
plant.GetFrameByName(self.end_effector_frame_name, arm),
X_ee,
arm,
)
)
return arm, end_effector_frame
[docs]
def add_gripper_controller(self):
"""
*Description*
This function adds a controller for the gripper to the station.
It is optional and may not be necessary for all stations.
"""
# Setup
# Create gripper controller
self.gripper_controller = self.builder.AddSystem(
GripperController(self.gripper_type)
)
self.gripper_controller.set_name(f"{self.get_name()}_gripper_controller")
[docs]
def AddArm(self):
"""
*Description*
Add the UR10e arm to the system.
"""
# Setup
arm_urdf_path = str(
impresources.files(robots) / "models/ur/ur10e.urdf",
)
arm_urdf = DrakeReadyURDFConverter(arm_urdf_path).convert_urdf()
arm_urdf = str(arm_urdf)
self.end_effector_frame_name = "tool0"
# Create transform for EE
self.X_ee = RigidTransform()
self.X_ee.set_translation([0, 0, 0.13])
# Add the arm to the plant using our convenience function
self.arm, self.arm_ee_frame = self.add_arm_to_plant_with_ee_frame(
arm_urdf, self.plant,
X_ee=self.X_ee,
)
# Add the arm to the controller plant using our convenience function
self.controller_arm, self.controller_arm_ee_frame = self.add_arm_to_plant_with_ee_frame(
arm_urdf, self.controller_plant,
X_ee=self.X_ee,
)
[docs]
def Add2f85Gripper(self):
"""
*Description*
Add the Robotiq 2F-85 gripper to the system. The arm must be added first.
"""
# Setup
# Add a gripper with actuation to the full simulated plant
gripper_urdf_path = str(
impresources.files(robots) / "models/robotiq/2f_85_gripper/urdf/robotiq_2f_85.urdf"
)
self.gripper = Parser(plant=self.plant).AddModels(gripper_urdf_path)[0]
X_grip = RigidTransform()
X_grip.set_rotation(
RotationMatrix(RollPitchYaw([0.0, 0.0, np.pi / 2]))
)
self.plant.WeldFrames(
self.plant.GetFrameByName("tool0", self.arm),
self.plant.GetFrameByName("robotiq_arg2f_base_link", self.gripper),
X_grip,
)
# # Add a gripper without actuation to the controller plant
# gripper_static_urdf = str(
# impresources.files(robots) / "models/robotiq/2f_85_gripper-no-mimic/urdf/robotiq_2f_85_static.urdf"
# )
# static_gripper = Parser(plant=self.controller_plant).AddModels(
# gripper_static_urdf
# )[0]
# self.controller_plant.WeldFrames(
# self.controller_plant.GetFrameByName("tool0", self.controller_arm),
# self.controller_plant.GetFrameByName("robotiq_arg2f_base_link", static_gripper),
# X_grip)
def ConnectToMeshcatVisualizer(self, port=None):
self.meshcat = Meshcat(port)
m = MeshcatVisualizer(self.meshcat)
m.AddToBuilder(self.builder, self.scene_graph, self.meshcat)
print("Open %s in a browser to view the meshcat visualizer." % self.meshcat.web_url())
[docs]
def CreateArmControllerAndConnect(self) -> InverseDynamicsController:
"""
Description
-----------
This function creates a Cartesian arm controller and connects it to the rest of the system.
:return:
"""
# Setup
kp = 1.0e3 * np.ones((self.plant.num_actuated_dofs(self.arm),))
kd = 2.0 * np.sqrt(kp)
# Create InverseDynamicsController
self.controller = self.builder.AddSystem(
InverseDynamicsController(
self.controller_plant,
kp=kp,
ki=0.0*kp,
kd=kd,
has_reference_acceleration=False,
)
)
# Torques from controller go to the "true" plant
self.builder.Connect(
self.controller.get_output_port_control(),
self.plant.get_actuation_input_port(self.arm),
)
self.builder.Connect(
self.controller.get_output_port_control(),
self.controller_plant.get_actuation_input_port(self.controller_arm),
)
# Connect plant's state to controller
self.builder.Connect(
self.plant.get_state_output_port(self.arm),
self.controller.get_input_port_estimated_state(),
)
return self.controller
[docs]
def CreateGripperControllerAndConnect(self):
"""
*Description*
Connects the gripper controller to:
- The Gripper Actuators (done via the plant of the station)
and exports some of the gripper controller's inputs and outputs
(so that they become the inputs and outputs of the station).
"""
# Setup
gripper_controller = self.gripper_controller
# Export the inputs of the gripper controller to the diagram
self.builder.ExportInput(
gripper_controller.GetInputPort("gripper_target"),
"gripper_target",
)
self.builder.ExportInput(
gripper_controller.GetInputPort("gripper_target_type"),
"gripper_target_type",
)
# Connect gripper controller to the plant
self.builder.Connect(
self.plant.get_state_output_port(self.gripper),
gripper_controller.GetInputPort("gripper_state"),
)
self.builder.Connect(
gripper_controller.GetOutputPort("applied_gripper_torque"),
self.plant.get_actuation_input_port(self.gripper),
)
# Send gripper position and velocity as an output
# self.builder.ExportOutput(
# gripper_controller.GetOutputPort("measured_gripper_position"),
# "measured_gripper_position",
# )
# self.builder.ExportOutput(
# gripper_controller.GetOutputPort("measured_gripper_velocity"),
# "measured_gripper_velocity",
# )
[docs]
def create_plants_and_scene_graph(self, time_step: float = 0.001):
"""
*Description*
This function creates the plant and scene graph for the diagram.
*Parameters*
time_step : float
The time step for the MultibodyPlant.
"""
# Setup
# Algorithm
self.plant, self.scene_graph = AddMultibodyPlantSceneGraph(
builder=self.builder,
time_step=time_step,
)
self.plant.set_name(f"{self.get_name()}_Plant")
self.scene_graph.set_name(f"{self.get_name()}_SceneGraph")
# Create plant for controller
self.controller_plant = self.builder.AddSystem(
MultibodyPlant(time_step=time_step)
)
self.controller_plant.set_name(f"{self.get_name()}_ControllerPlant")
[docs]
def ExportArmControllerPorts(self):
"""
*Description*
This function creates the input and output ports for the diagram
that are related to the arm controller.
"""
# Setup
controller = self.controller
controller_plant = controller.get_multibody_plant_for_control()
# Create commanded arm torque output
self.builder.ExportOutput(
self.controller.get_output_port_control(),
"ur10e_arm_torque_commanded",
)
# Create desired position input
self.ExportArmPositionInputPort()
# Create Multiplexer to extract the arm position and velocity
demux = self.builder.AddSystem(
Demultiplexer(
self.plant.num_multibody_states(self.arm),
self.plant.num_positions(self.arm),
),
)
demux.set_name("UR10e_Station_Demux")
self.builder.Connect(
self.plant.get_state_output_port(self.arm),
demux.get_input_port(0),
)
self.builder.ExportOutput(
demux.get_output_port(0),
"ur10e_arm_position_measured",
)
self.builder.ExportOutput(
demux.get_output_port(1),
"ur10e_arm_velocity_measured",
)
# Compute and output end-effector wrenches based on measured joint torques
wrench_calculator = self.builder.AddSystem(
EndEffectorWrenchCalculator(
controller_plant,
controller_plant.GetFrameByName("end_effector"),
),
)
wrench_calculator.set_name("wrench_calculator")
self.builder.Connect(
demux.get_output_port(0),
wrench_calculator.GetInputPort("joint_positions"))
self.builder.Connect(
demux.get_output_port(1),
wrench_calculator.GetInputPort("joint_velocities"))
self.builder.Connect(
self.controller.get_output_port_control(),
wrench_calculator.GetInputPort("joint_torques"))
self.builder.ExportOutput(
wrench_calculator.get_output_port(),
"measured_ee_wrench",
)
[docs]
def ExportGripperStatePort(self):
"""
*Description*
This function exports the gripper's state
as a port of the diagram.
"""
# Setup
# Export the state of the gripper
self.builder.ExportOutput(
self.plant.get_state_output_port(self.gripper),
"gripper_state",
)
[docs]
def Finalize(self):
"""
*Description*
This finalizes the diagram by:
- Finalizing all plants in the diagram
- Setting up the visualizer (if desired)
- Creating and connecting the arm controller
- Creating and connecting the gripper controller (if desired)
- Building the diagram
"""
# Setup
# Finalize all plants
self.plant.Finalize()
self.controller_plant.Finalize()
# # Set up the scene graph
# self.builder.Connect(
# self.scene_graph.get_query_output_port(),
# self.plant.get_geometry_query_input_port()
# )
# self.builder.Connect(
# self.plant.get_geometry_pose_output_port(),
# self.scene_graph.get_source_pose_port(self.plant.get_source_id())
# )
# Connect Meshcat, if desired
if self.use_meshcat():
self.ConnectToMeshcatVisualizer(port=self.meshcat_port_number)
# Create Arm and Gripper Controllers
self.CreateArmControllerAndConnect()
self.ExportArmControllerPorts()
if self.gripper_type != GripperType.NoGripper:
self.CreateGripperControllerAndConnect()
self.ExportGripperStatePort()
# Build the diagram
self.builder.BuildInto(self)
[docs]
def use_meshcat(self) -> bool:
"""
*Description*
This function returns whether or not to use meshcat for visualization.
:return: (bool) Whether or not to use meshcat for visualization.
"""
return self.meshcat_port_number is not None