def __init__(self, fname, idn=0):
# Always run super init first
super().__init__(fname, idn=idn)
# key: gripper name and value: gripper model
self.grippers = OrderedDict()
# Grab hand's offset from final robot link (string -> np.array -> elements [1, 2, 3, 0] (x, y, z, w))
# Different case based on whether we're dealing with single or bimanual armed robot
if self.arm_type == "single":
hand_element = find_elements(root=self.root,
tags="body",
attribs={"name": self.eef_name},
return_first=True)
self.hand_rotation_offset = string_to_array(
hand_element.get("quat", "1 0 0 0"))[[1, 2, 3, 0]]
else: # "bimanual" case
self.hand_rotation_offset = {}
for arm in ("right", "left"):
hand_element = find_elements(
root=self.root,
tags="body",
attribs={"name": self.eef_name[arm]},
return_first=True)
self.hand_rotation_offset[arm] = string_to_array(
hand_element.get("quat", "1 0 0 0"))[[1, 2, 3, 0]]
# Get camera names for this robot
self.cameras = self.get_element_names(self.worldbody, "camera")
def set_camera(self, camera_name, pos, quat, camera_attribs=None):
"""
Sets a camera with @camera_name. If the camera already exists, then this overwrites its pos and quat values.
Args:
camera_name (str): Camera name to search for / create
pos (3-array): (x,y,z) coordinates of camera in world frame
quat (4-array): (w,x,y,z) quaternion of camera in world frame
camera_attribs (dict): If specified, should be additional keyword-mapped attributes for this camera.
See http://www.mujoco.org/book/XMLreference.html#camera for exact attribute specifications.
"""
# Determine if camera already exists
camera = find_elements(root=self.worldbody,
tags="camera",
attribs={"name": camera_name},
return_first=True)
# Compose attributes
if camera_attribs is None:
camera_attribs = {}
camera_attribs["pos"] = array_to_string(pos)
camera_attribs["quat"] = array_to_string(quat)
if camera is None:
# If camera doesn't exist, then add a new camera with the specified attributes
self.worldbody.append(
new_element(tag="camera", name=camera_name, **camera_attribs))
else:
# Otherwise, we edit all specified attributes in that camera
for attrib, value in camera_attribs.items():
camera.set(attrib, value)
def _append_joints(self, root, body_name=None, joint_specs="default"):
"""
Appends all joints as specified by @joint_specs to @body.
Args:
root (ET.Element): Top-level element to iteratively search through for @body_name
body_name (None or str): Name of the body to append the joints to.
None defaults to "root" (top-level body)
joint_specs (str or list): List of joint specifications to add to the specified body, or
"default", which results in a single free joint
"""
# Standardize inputs
if body_name is None:
body_name = "root"
if joint_specs == "default":
joint_specs = [self.get_joint_attrib_template()]
for i, joint_spec in enumerate(joint_specs):
if "name" not in joint_spec:
joint_spec["name"] = f"{body_name}_joint{i}"
# Search for body and make sure it exists
body = find_elements(root=root,
tags="body",
attribs={"name": body_name},
return_first=True)
assert body is not None, "Could not find body with name: {}".format(
body_name)
# Add joint(s) to this body
for joint_spec in joint_specs:
body.append(new_joint(**joint_spec))
def __init__(self):
super().__init__(xml_path_completion("base.xml"))
# Modify the simulation timestep to be the requested value
options = find_elements(root=self.root,
tags="option",
attribs=None,
return_first=True)
options.set("timestep", convert_to_string(macros.SIMULATION_TIMESTEP))
def merge_assets(self, other):
"""
Merges @other's assets in a custom logic.
Args:
other (MujocoXML or MujocoObject): other xml file whose assets will be merged into this one
"""
for asset in other.asset:
if find_elements(root=self.asset, tags=asset.tag,
attribs={"name": asset.get("name")}, return_first=True) is None:
self.asset.append(asset)
def _set_door_damping(self, damping):
"""
Helper function to override the door friction directly in the XML
Args:
damping (float): damping parameter to override the ones specified in the XML
"""
hinge = find_elements(root=self.worldbody,
tags="joint",
attribs={"name": self.hinge_joint},
return_first=True)
hinge.set("damping", array_to_string(np.array([damping])))
def _set_door_friction(self, friction):
"""
Helper function to override the door friction directly in the XML
Args:
friction (3-tuple of float): friction parameters to override the ones specified in the XML
"""
hinge = find_elements(root=self.worldbody,
tags="joint",
attribs={"name": self.hinge_joint},
return_first=True)
hinge.set("frictionloss", array_to_string(np.array([friction])))
def modify_xml_for_camera_movement(xml, camera_name):
"""
Cameras in mujoco are 'fixed', so they can't be moved by default.
Although it's possible to hack position movement, rotation movement
does not work. An alternative is to attach a camera to a body element,
and move the body.
This function modifies the camera with name @camera_name in the xml
by attaching it to a body element that we can then manipulate. In this
way, we can move the camera by moving the body.
See http://www.mujoco.org/forum/index.php?threads/move-camera.2201/ for
further details.
xml (str): Mujoco sim XML file as a string
camera_name (str): Name of camera to tune
"""
tree = ET.fromstring(xml)
# find the correct camera
camera_elem = None
cameras = find_elements(root=tree, tags="camera", return_first=False)
for camera in cameras:
if camera.get("name") == camera_name:
camera_elem = camera
break
assert camera_elem is not None, "No valid camera name found, options are: {}"\
.format([camera.get("name") for camera in cameras])
# Find parent element of the camera element
parent = find_parent(root=tree, child=camera_elem)
assert parent is not None
# add camera body
cam_body = ET.SubElement(parent, "body")
cam_body.set("name", "cameramover")
cam_body.set("pos", camera_elem.get("pos"))
cam_body.set("quat", camera_elem.get("quat"))
new_camera = ET.SubElement(cam_body, "camera")
new_camera.set("mode", "fixed")
new_camera.set("name", camera_elem.get("name"))
new_camera.set("pos", "0 0 0")
# Also need to define inertia
inertial = ET.SubElement(cam_body, "inertial")
inertial.set("diaginertia", "1e-08 1e-08 1e-08")
inertial.set("mass", "1e-08")
inertial.set("pos", "0 0 0")
# remove old camera element
parent.remove(camera_elem)
return ET.tostring(tree, encoding="utf8").decode("utf8")
def configure_location(self):
"""Configures correct locations for this arena"""
# Run superclass first
super().configure_location()
# Define start position for drawing the line
pos = self.sample_start_pos()
# Define dirt material for markers
tex_attrib = {
"type": "cube",
}
mat_attrib = {
"texrepeat": "1 1",
"specular": "0.0",
"shininess": "0.0",
}
dirt = CustomMaterial(
texture="Dirt",
tex_name="dirt",
mat_name="dirt_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
shared=True,
)
# Define line(s) drawn on table
for i in range(self.num_markers):
# If we're using two clusters, we resample the starting position and direction at the halfway point
if self.two_clusters and i == int(np.floor(self.num_markers / 2)):
pos = self.sample_start_pos()
marker_name = f'contact{i}'
marker = CylinderObject(
name=marker_name,
size=[self.line_width / 2, 0.001],
rgba=[1, 1, 1, 1],
material=dirt,
obj_type="visual",
joints=None,
)
# Manually add this object to the arena xml
self.merge_assets(marker)
table = find_elements(root=self.worldbody, tags="body", attribs={"name": "table"}, return_first=True)
table.append(marker.get_obj())
# Add this marker to our saved list of all markers
self.markers.append(marker)
# Add to the current dirt path
pos = self.sample_path_pos(pos)
def _append_object(self, root, obj, parent_name=None, pos=None, quat=None):
"""
Helper function to add pre-generated object @obj to the body with name @parent_name
Args:
root (ET.Element): Top-level element to iteratively search through for @parent_name to add @obj to
obj (MujocoObject): Object to append to the body specified by @parent_name
parent_name (None or str): Body name to search for in @root to append @obj to.
None defaults to "root" (top-level body)
pos (None or 3-array): (x,y,z) relative offset from parent body when appending @obj.
None defaults to (0,0,0)
quat (None or 4-array) (w,x,y,z) relative quaternion rotation from parent body when appending @obj.
None defaults to (1,0,0,0)
"""
# Set defaults if any are None
if parent_name is None:
parent_name = "root"
if pos is None:
pos = np.zeros(3)
if quat is None:
quat = np.array([1, 0, 0, 0])
# First, find parent body
parent = find_elements(root=root,
tags="body",
attribs={"name": parent_name},
return_first=True)
assert parent is not None, "Could not find parent body with name: {}".format(
parent_name)
# Get the object xml element tree, remove its top-level joints, and modify its top-level pos / quat
child = obj.get_obj()
self._remove_joints(child)
child.set("pos", array_to_string(pos))
child.set("quat", array_to_string(quat))
# Add this object and its assets to this composite object
self.merge_assets(other=obj)
parent.append(child)
# Update geometric properties for this composite object
obj_abs_pos = self._object_absolute_positions[parent_name] + np.array(
pos)
self._object_absolute_positions[obj.root_body] = obj_abs_pos
self._top = max(self._top, obj_abs_pos[2] + obj.top_offset[2])
self._bottom = min(self._bottom, obj_abs_pos[2] + obj.bottom_offset[2])
self._horizontal = max(self._horizontal,
max(obj_abs_pos[:2]) + obj.horizontal_radius)
def merge(self, others, merge_body="default"):
"""
Default merge method.
Args:
others (MujocoXML or list of MujocoXML): other xmls to merge into this one
raises XML error if @others is not a MujocoXML instance.
merges <worldbody/>, <actuator/> and <asset/> of @others into @self
merge_body (None or str): If set, will merge child bodies of @others. Default is "default", which
corresponds to the root worldbody for this XML. Otherwise, should be an existing body name
that exists in this XML. None results in no merging of @other's bodies in its worldbody.
Raises:
XMLError: [Invalid XML instance]
"""
if type(others) is not list:
others = [others]
for idx, other in enumerate(others):
if not isinstance(other, MujocoXML):
raise XMLError("{} is not a MujocoXML instance.".format(
type(other)))
if merge_body is not None:
root = (self.worldbody if merge_body == "default" else
find_elements(root=self.worldbody,
tags="body",
attribs={"name": merge_body},
return_first=True))
for body in other.worldbody:
root.append(body)
self.merge_assets(other)
for one_actuator in other.actuator:
self.actuator.append(one_actuator)
for one_sensor in other.sensor:
self.sensor.append(one_sensor)
for one_tendon in other.tendon:
self.tendon.append(one_tendon)
for one_equality in other.equality:
self.equality.append(one_equality)
for one_contact in other.contact:
self.contact.append(one_contact)
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose(s) accordingly
if self.env_configuration == "bimanual":
xpos = self.robots[0].robot_model.base_xpos_offset["empty"]
self.robots[0].robot_model.set_base_xpos(xpos)
else:
if self.env_configuration == "single-arm-opposed":
# Set up robots facing towards each other by rotating them from their default position
for robot, rotation in zip(self.robots,
(np.pi / 2, -np.pi / 2)):
xpos = robot.robot_model.base_xpos_offset["empty"]
rot = np.array((0, 0, rotation))
xpos = T.euler2mat(rot) @ np.array(xpos)
robot.robot_model.set_base_xpos(xpos)
robot.robot_model.set_base_ori(rot)
else: # "single-arm-parallel" configuration setting
# Set up robots parallel to each other but offset from the center
for robot, offset in zip(self.robots, (-0.25, 0.25)):
xpos = robot.robot_model.base_xpos_offset["empty"]
xpos = np.array(xpos) + np.array((0, offset, 0))
robot.robot_model.set_base_xpos(xpos)
# Add arena and robot
mujoco_arena = EmptyArena()
# Arena always gets set to zero origin
mujoco_arena.set_origin([0, 0, 0])
# Modify default agentview camera
mujoco_arena.set_camera(camera_name="agentview",
pos=[
1.0666432116509934, 1.4903257668114777e-08,
2.0563394967349096
],
quat=[
0.6530979871749878, 0.27104058861732483,
0.27104055881500244, 0.6530978679656982
])
# initialize objects of interest
self.hole = PlateWithHoleObject(name="hole")
tex_attrib = {
"type": "cube",
}
mat_attrib = {
"texrepeat": "1 1",
"specular": "0.4",
"shininess": "0.1",
}
greenwood = CustomMaterial(
texture="WoodGreen",
tex_name="greenwood",
mat_name="greenwood_mat",
tex_attrib=tex_attrib,
mat_attrib=mat_attrib,
)
self.peg = CylinderObject(
name="peg",
size_min=(self.peg_radius[0], self.peg_length),
size_max=(self.peg_radius[1], self.peg_length),
material=greenwood,
rgba=[0, 1, 0, 1],
joints=None,
)
# Load hole object
hole_obj = self.hole.get_obj()
hole_obj.set("quat", "0 0 0.707 0.707")
hole_obj.set("pos", "0.11 0 0.17")
# Load peg object
peg_obj = self.peg.get_obj()
peg_obj.set("pos", array_to_string((0, 0, self.peg_length)))
# Append appropriate objects to arms
if self.env_configuration == "bimanual":
r_eef, l_eef = [
self.robots[0].robot_model.eef_name[arm]
for arm in self.robots[0].arms
]
r_model, l_model = [
self.robots[0].robot_model, self.robots[0].robot_model
]
else:
r_eef, l_eef = [
robot.robot_model.eef_name for robot in self.robots
]
r_model, l_model = [
self.robots[0].robot_model, self.robots[1].robot_model
]
r_body = find_elements(root=r_model.worldbody,
tags="body",
attribs={"name": r_eef},
return_first=True)
l_body = find_elements(root=l_model.worldbody,
tags="body",
attribs={"name": l_eef},
return_first=True)
r_body.append(peg_obj)
l_body.append(hole_obj)
# task includes arena, robot, and objects of interest
# We don't add peg and hole directly since they were already appended to the robots
self.model = ManipulationTask(
mujoco_arena=mujoco_arena,
mujoco_robots=[robot.robot_model for robot in self.robots],
)
# Make sure to add relevant assets from peg and hole objects
self.model.merge_assets(self.hole)
self.model.merge_assets(self.peg)
