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 _get_collision(self, site=False, ob_type="box"):
main_body = ET.Element("body")
main_body.set("name", self.name)
template = self.get_collision_attrib_template()
template["name"] = self.name
template["type"] = ob_type
if self.material == "default":
template["rgba"] = "0.5 0.5 0.5 1" # mujoco default
template["material"] = "{}_mat".format(self.name)
elif self.material is not None:
template["material"] = self.material.mat_attrib["name"]
else:
template["rgba"] = array_to_string(self.rgba)
template["size"] = array_to_string(self.size)
template["density"] = str(self.density)
template["friction"] = array_to_string(self.friction)
template["solref"] = array_to_string(self.solref)
template["solimp"] = array_to_string(self.solimp)
main_body.append(ET.Element("geom", attrib=template))
if site:
# add a site as well
template = self.get_site_attrib_template()
template["name"] = self.name
main_body.append(ET.Element("site", attrib=template))
return main_body
def place_objects(self):
"""Places objects randomly until no collisions or max iterations hit."""
pos_arr, quat_arr = self.initializer.sample(self.push_object_idx)
for i in range(len(self.xml_objects)):
self.xml_objects[i].set("pos", array_to_string(pos_arr[i]))
self.xml_objects[i].set("quat", array_to_string(quat_arr[i]))
def configure_location(self):
"""Configures correct locations for this arena"""
# Run superclass first
super().configure_location()
# Determine peg friction
friction = max(0.001, np.random.normal(self.table_friction[0], self.table_friction_std))
# Grab reference to the table body in the xml
table_subtree = self.worldbody.find(".//body[@name='{}']".format("table"))
# 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,
)
# Define line(s) drawn on table
for i in range(self.num_sensors):
# 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_sensors / 2)):
pos = self.sample_start_pos()
square_name2 = 'contact_'+str(i)
square2 = CylinderObject(
name=square_name2,
size=[self.line_width / 2, 0.001],
rgba=[1, 1, 1, 1],
density=1,
material=dirt,
friction=friction,
)
self.merge_asset(square2)
visual_c = square2.get_visual(site=True)
visual_c.set("pos", array_to_string([pos[0], pos[1], self.table_half_size[2]]))
visual_c.find("site").set("pos", [0, 0, 0.005])
visual_c.find("site").set("rgba", array_to_string([0, 0, 0, 0]))
table_subtree.append(visual_c)
sensor_name = square_name2 + "_sensor"
sensor_site_name = square_name2
self.sensor_names += [sensor_name]
self.sensor_site_names[sensor_name] = sensor_site_name
# Add to the current dirt path
pos = self.sample_path_pos(pos)
def place_objects_not_random(self):
"""Places objects randomly until no collisions or max iterations hit."""
#pos_arr = [np.array([ 0.53522776, -0.0287869 , 0.82162434])]
#quat_arr = [[0.8775825618903728, 0, 0, 0.479425538604203]]
pos_arr = self.box_pos_array
quat_arr = self.box_quat_array
for i in range(len(self.objects)):
self.objects[i].set("pos", array_to_string(pos_arr[i]))
self.objects[i].set("quat", array_to_string(quat_arr[i]))
def place_objects_plus_fixed(self):
"""Places objects randomly until no collisions or max iterations hit.
and the object in use to a defined location
"""
placed_objects = []
index = 0
# place objects by rejection sampling
for _, obj_mjcf in self.mujoco_objects.items():
horizontal_radius = obj_mjcf.get_horizontal_radius()
bottom_offset = obj_mjcf.get_bottom_offset()
success = False
# Fixing position of defined object
if index == self.obj_id:
pos = self.obj_fixed_loc[0]
placed_objects.append((pos, horizontal_radius))
self.objects[index].set("pos", array_to_string(pos))
quat = self.obj_fixed_loc[1]
self.objects[index].set("quat", array_to_string(quat))
success = True
continue
for _ in range(5000): # 5000 retries
bin_x_half = self.bin_size[0] / 2 - horizontal_radius - 0.05
bin_y_half = self.bin_size[1] / 2 - horizontal_radius - 0.05
object_x = np.random.uniform(high=bin_x_half, low=-bin_x_half)
object_y = np.random.uniform(high=bin_y_half, low=-bin_y_half)
# make sure objects do not overlap
object_xy = np.array([object_x, object_y, 0])
pos = self.bin_offset - bottom_offset + object_xy
location_valid = True
for pos2, r in placed_objects:
dist = np.linalg.norm(pos[:2] - pos2[:2], np.inf)
if dist <= r + horizontal_radius:
location_valid = False
break
# place the object
if location_valid:
# add object to the position
placed_objects.append((pos, horizontal_radius))
self.objects[index].set("pos", array_to_string(pos))
# random z-rotation
quat = self.sample_quat()
self.objects[index].set("quat", array_to_string(quat))
success = True
break
# raise error if all objects cannot be placed after maximum retries
if not success:
raise RandomizationError("Cannot place all objects in the bins")
index += 1
def move_objects_random(self):
"""Places objects randomly until no collisions or max iterations hit."""
placed_objects = []
index = 0
count = 0
# place objects by rejection sampling
for _, obj_mjcf in self.mujoco_objects.items():
horizontal_radius = obj_mjcf.get_horizontal_radius()
bottom_offset = obj_mjcf.get_bottom_offset()
success = False
for _ in range(5000): # 5000 retries
bin_x_half = 0.1 - horizontal_radius - 0.01
bin_y_half = 0.15 - horizontal_radius - 0.01
object_x = np.random.uniform(high=bin_x_half, low=-bin_x_half)
object_y = np.random.uniform(high=bin_y_half, low=-bin_y_half)
# make sure objects do not overlap
object_xy = np.array([object_x, object_y, 0])
pos = self.bin2_offset - bottom_offset + object_xy
location_valid = True
for pos2, r in placed_objects:
dist = np.linalg.norm(pos[:2] - pos2[:2], np.inf)
if dist <= r + horizontal_radius:
location_valid = False
break
# place the object
if location_valid:
# add object to the position
placed_objects.append((pos, horizontal_radius))
self.objects[index].set("pos", array_to_string(pos))
# random z-rotation
quat = self.sample_quat()
self.objects[index].set("quat", array_to_string(quat))
success = True
# time.sleep(1)
break
if success:
count += 1
# raise error if all objects cannot be placed after maximum retries
# if not success:
# raise RandomizationError("Cannot place all objects in the bins")
index += 1
print("Placed %d objects in bin" % count)
def drop_cube_on_body_demo():
world = MujocoWorldBase()
arena = EmptyArena()
arena.set_origin([0, 0, 0])
world.merge(arena)
soft_torso = SoftTorsoObject()
obj = soft_torso.get_collision()
box = BoxObject()
box_obj = box.get_collision()
obj.append(new_joint(name='soft_torso_free_joint', type='free'))
box_obj.append(new_joint(name='box_free_joint', type='free'))
world.merge_asset(soft_torso)
world.worldbody.append(obj)
world.worldbody.append(box_obj)
# Place torso on ground
collision_soft_torso = world.worldbody.find("./body")
collision_soft_torso.set("pos", array_to_string(np.array([-0.1, 0, 0.1])))
model = world.get_model(mode="mujoco_py")
sim = MjSim(model)
viewer = MjViewer(sim)
for _ in range(10000):
sim.step()
viewer.render()
def place_visual(self):
"""Places visual objects randomly until no collisions or max iterations hit."""
index = 0
bin_pos = string_to_array(self.bin2_body.get("pos"))
bin_size = self.bin_size
for _, obj_mjcf in self.visual_objects:
bin_x_low = bin_pos[0]
bin_y_low = bin_pos[1]
if index == 0 or index == 2:
bin_x_low -= bin_size[0] / 2
if index < 2:
bin_y_low -= bin_size[1] / 2
bin_x_high = bin_x_low + bin_size[0] / 2
bin_y_high = bin_y_low + bin_size[1] / 2
bottom_offset = obj_mjcf.get_bottom_offset()
bin_range = [bin_x_low + bin_x_high, bin_y_low + bin_y_high, 2 * bin_pos[2]]
bin_center = np.array(bin_range) / 2.0
pos = bin_center - bottom_offset
self.visual_obj_mjcf[index].set("pos", array_to_string(pos))
index += 1
def set_origin(self, offset):
"""Applies a constant offset to all objects."""
offset = np.array(offset)
for node in self.worldbody.findall("./*[@pos]"):
cur_pos = string_to_array(node.get("pos"))
new_pos = cur_pos + offset
node.set("pos", array_to_string(new_pos))
def add_mount(self, mount):
"""
Mounts @mount to arm.
Throws error if robot already has a mount or if mount type is incorrect.
Args:
mount (MountModel): mount MJCF model
Raises:
ValueError: [mount already added]
"""
if self.mount is not None:
raise ValueError("Mount already added for this robot!")
# First adjust mount's base position
offset = self.base_offset - mount.top_offset
mount._elements["root_body"].set("pos", array_to_string(offset))
self.merge(mount, merge_body=self.root_body)
self.mount = mount
# Update cameras in this model
self.cameras = self.get_element_names(self.worldbody, "camera")
def place_objects(self):
"""
Places objects randomly on table until no collisions or max iterations hit.
"""
pos_arr, _ = self.initializer.sample()
for i in range(len(self.objects_on_table)):
self.objects_on_table[i].set("pos", array_to_string(pos_arr[i]))
return pos_arr
def _get_cube_material(self):
from robosuite.utils.mjcf_utils import array_to_string
rgba = (1, 0, 0, 1)
cube_material = CustomMaterial(
texture=rgba,
tex_name="solid",
mat_name="solid_mat",
)
cube_material.tex_attrib.pop('file')
cube_material.tex_attrib["type"] = "cube"
cube_material.tex_attrib["builtin"] = "flat"
cube_material.tex_attrib["rgb1"] = array_to_string(rgba[:3])
cube_material.tex_attrib["rgb2"] = array_to_string(rgba[:3])
cube_material.tex_attrib["width"] = "100"
cube_material.tex_attrib["height"] = "100"
return cube_material
def get_collision_attrib_template():
"""
Generates template with collision attributes for a given geom
Returns:
dict: Initial template with `'pos'` and `'group'` already specified
"""
return {"group": "0", "rgba": array_to_string(OBJECT_COLLISION_COLOR)}
def set_joint_damping(self,
damping=np.array((0.1, 0.1, 0.1, 0.1, 0.1, 0.01))):
"""Set joint damping """
body = self._base_body
for i in range(len(self._link_body)):
body = body.find("./body[@name='{}']".format(self._link_body[i]))
joint = body.find("./joint[@name='{}']".format(self._joints[i]))
joint.set("damping", array_to_string(np.array([damping[i]])))
def set_joint_frictionloss(self,
friction=np.array(
(0.1, 0.1, 0.1, 0.1, 0.1, 0.01))):
"""Set joint friction loss (static friction)"""
body = self._base_body
for i in range(len(self._link_body)):
body = body.find("./body[@name='{}']".format(self._link_body[i]))
joint = body.find("./joint[@name='{}']".format(self._joints[i]))
joint.set("frictionloss", array_to_string(np.array([friction[i]])))
def set_base_xpos(self, pos):
"""
Places the robot on position @pos.
Args:
pos (3-array): (x,y,z) position to place robot base
"""
node = self.worldbody.find("./body[@name='{}']".format(self._root_))
node.set("pos", array_to_string(pos - self.bottom_offset))
def set_base_xpos(self, pos):
"""
Places the robot on position @pos.
Args:
pos (3-array): (x,y,z) position to place robot base
"""
self._elements["root_body"].set(
"pos", array_to_string(pos - self.bottom_offset))
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 set_base_ori(self, rot):
"""
Rotates robot by rotation @rot from its original orientation.
Args:
rot (3-array): (r,p,y) euler angles specifying the orientation for the robot base
"""
# xml quat assumes w,x,y,z so we need to convert to this format from outputted x,y,z,w format from fcn
rot = mat2quat(euler2mat(rot))[[3, 0, 1, 2]]
self._elements["root_body"].set("quat", array_to_string(rot))
def set_base_ori(self, rot):
"""
Rotates robot by rotation @rot from its original orientation.
Args:
rot (3-array): (r,p,y) euler angles specifying the orientation for the robot base
"""
node = self.worldbody.find("./body[@name='{}']".format(self._root_))
# xml quat assumes w,x,y,z so we need to convert to this format from outputted x,y,z,w format from fcn
rot = mat2quat(euler2mat(rot))[[3, 0, 1, 2]]
node.set("quat", array_to_string(rot))
def configure_location(self):
"""Configures correct locations for this arena"""
self.bottom_pos = np.array([0, 0, 0])
self.floor.set("pos", array_to_string(self.bottom_pos))
self.center_pos = self.bottom_pos + np.array(
[0, 0, -self.table_half_size[2]]) + self.table_offset
self.table_body.set("pos", array_to_string(self.center_pos))
self.table_collision.set("size", array_to_string(self.table_half_size))
self.table_collision.set("friction",
array_to_string(self.table_friction))
self.table_visual.set("size", array_to_string(self.table_half_size))
self.table_top.set(
"pos", array_to_string(np.array([0, 0, self.table_half_size[2]])))
# If we're not using legs, set their size to 0
if not self.has_legs:
for leg in self.table_legs_visual:
leg.set("rgba", array_to_string([1, 0, 0, 0]))
leg.set("size", array_to_string([0.0001, 0.0001]))
else:
# Otherwise, set leg locations appropriately
delta_x = [0.1, -0.1, -0.1, 0.1]
delta_y = [0.1, 0.1, -0.1, -0.1]
for leg, dx, dy in zip(self.table_legs_visual, delta_x, delta_y):
# If x-length of table is less than a certain length, place leg in the middle between ends
# Otherwise we place it near the edge
x = 0
if self.table_half_size[0] > abs(dx * 2.0):
x += np.sign(dx) * self.table_half_size[0] - dx
# Repeat the same process for y
y = 0
if self.table_half_size[1] > abs(dy * 2.0):
y += np.sign(dy) * self.table_half_size[1] - dy
# Get z value
z = (self.table_offset[2] - self.table_half_size[2]) / 2.0
# Set leg position
leg.set("pos", array_to_string([x, y, -z]))
# Set leg size
leg.set("size", array_to_string([0.025, z]))
def _get_object_subtree_(self, ob_type="box"):
# Create element tree
obj = new_body(name="main")
# Get base element attributes
element_attr = {
"name": "g0",
"type": ob_type,
"size": array_to_string(self.size)
}
# Add collision geom if necessary
if self.obj_type in {"collision", "all"}:
col_element_attr = deepcopy(element_attr)
col_element_attr.update(self.get_collision_attrib_template())
col_element_attr["density"] = str(self.density)
col_element_attr["friction"] = array_to_string(self.friction)
col_element_attr["solref"] = array_to_string(self.solref)
col_element_attr["solimp"] = array_to_string(self.solimp)
obj.append(new_geom(**col_element_attr))
# Add visual geom if necessary
if self.obj_type in {"visual", "all"}:
vis_element_attr = deepcopy(element_attr)
vis_element_attr.update(self.get_visual_attrib_template())
vis_element_attr["name"] += "_vis"
if self.material == "default":
vis_element_attr["rgba"] = "0.5 0.5 0.5 1" # mujoco default
vis_element_attr["material"] = "mat"
elif self.material is not None:
vis_element_attr["material"] = self.material.mat_attrib["name"]
else:
vis_element_attr["rgba"] = array_to_string(self.rgba)
obj.append(new_geom(**vis_element_attr))
# add joint(s)
for joint_spec in self.joint_specs:
obj.append(new_joint(**joint_spec))
# add a site as well
site_element_attr = self.get_site_attrib_template()
site_element_attr["name"] = "default_site"
obj.append(new_site(**site_element_attr))
return obj
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
"""
collision = self.worldbody.find("./body/body[@name='collision']")
node = collision.find("./body[@name='frame']")
node = node.find("./body[@name='door']")
hinge = node.find("./joint[@name='door_hinge']")
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 set_origin(self, offset):
"""
Applies a constant offset to all objects.
Args:
offset (3-tuple): (x,y,z) offset to apply to all nodes in this XML
"""
offset = np.array(offset)
for node in self.worldbody.findall("./*[@pos]"):
cur_pos = string_to_array(node.get("pos"))
new_pos = cur_pos + offset
node.set("pos", array_to_string(new_pos))
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
"""
collision = self.worldbody.find("./body/body[@name='collision']")
node = collision.find("./body[@name='frame']")
node = node.find("./body[@name='door']")
hinge = node.find("./joint[@name='door_hinge']")
hinge.set("frictionloss", array_to_string(np.array([friction])))
def set_goal_xpos(self, x_delta, y_delta):
""" Sets x,y position of goal site in door model with x and y offset from door center"""
door_center_site = self.worldbody.find(
"./body/body/body/site[@name='door_center']")
door_center_pos = string_to_array(door_center_site.get("pos"))
goal_site = self.worldbody.find("./body/body/body/site[@name='goal']")
goal_site.set(
"pos",
array_to_string([
door_center_pos[0] + x_delta, door_center_pos[1] + y_delta,
-1.0
]))
def set_stiffness(self, stiffness):
"""
Helper function to override the soft body's stiffness directly in the XML
Args:
stiffness (float, must be greater than zero): stiffness parameter to override the ones specified in the XML
"""
assert stiffness > 0, 'Damping must be greater than zero'
composite = self._get_composite_element()
solref_str = composite.get('solrefsmooth').split(' ')
damping = float(solref_str[1])
solref = np.array([-stiffness, damping])
composite.set('solrefsmooth', array_to_string(solref))
