def __init__(
self,
cylinder_radius=(0.015, 0.03),
cylinder_length=0.13,
use_object_obs=True,
reward_shaping=True,
**kwargs
):
"""
Args:
cylinder_radius (2-tuple): low and high limits of the (uniformly sampled)
radius of the cylinder
cylinder_length (float): length of the cylinder
use_object_obs (bool): if True, include object information in the observation.
reward_shaping (bool): if True, use dense rewards
Inherits the Baxter environment; refer to other parameters described there.
"""
# initialize objects of interest
self.hole = PlateWithHoleObject()
cylinder_radius = np.random.uniform(0.015, 0.03)
self.cylinder = CylinderObject(
size_min=(cylinder_radius, cylinder_length),
size_max=(cylinder_radius, cylinder_length),
)
self.mujoco_objects = OrderedDict()
# whether to use ground-truth object states
self.use_object_obs = use_object_obs
# reward configuration
self.reward_shaping = reward_shaping
super().__init__(gripper_left=None, gripper_right=None, **kwargs)
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
self.model = MujocoWorldBase()
self.mujoco_arena = EmptyArena()
if self.use_indicator_object:
self.mujoco_arena.add_pos_indicator()
self.model.merge(self.mujoco_arena)
for robot in self.robots:
self.model.merge(robot.robot_model)
# 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],
)
# Load hole object
self.hole_obj = self.hole.get_collision(site=True)
self.hole_obj.set("quat", "0 0 0.707 0.707")
self.hole_obj.set("pos", "0.11 0 0.17")
self.model.merge_asset(self.hole)
# Load peg object
self.peg_obj = self.peg.get_collision(site=True)
self.peg_obj.set("pos", array_to_string((0, 0, self.peg_length)))
self.model.merge_asset(self.peg)
# Depending on env configuration, append appropriate objects to arms
if self.env_configuration == "bimanual":
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[0].robot_model.eef_name["left"])).append(
self.hole_obj)
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[0].robot_model.eef_name["right"])).append(
self.peg_obj)
else:
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[1].robot_model.eef_name)).append(self.hole_obj)
self.model.worldbody.find(".//body[@name='{}']".format(
self.robots[0].robot_model.eef_name)).append(self.peg_obj)
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)
def _load_model(self):
"""
Loads an xml model, puts it in self.model
"""
super()._load_model()
# Adjust base pose accordingly
xpos = self.robots[0].robot_model.base_xpos_offset["table"](
self.table_full_size[0])
self.robots[0].robot_model.set_base_xpos(xpos)
# load model for table top workspace
mujoco_arena = TableArena(
table_full_size=self.table_full_size,
table_friction=self.table_friction,
table_offset=self.table_offset,
)
# Arena always gets set to zero origin
mujoco_arena.set_origin([0, 0, 0])
# Load hole object
# register object with the corresponding option (objectClass, name, xrange, yrange)
if self.task_configs['board'] == 'GMC_assembly':
self.register_object(my_object.GMC_Assembly_Object,
'plate',
xrange=[0, 0],
yrange=[0, 0])
if self.task_configs['board'] == 'GMC_plate':
self.register_object(my_object.GMC_Plate_Object,
'plate',
xrange=[0, 0],
yrange=[0, 0])
if self.task_configs['board'] == 'Square_hole_16mm':
self.register_object(my_object.square_hole_16mm,
'plate',
xrange=[0, 0],
yrange=[0, 0])
# Load peg object
if self.task_configs['peg'] == '16mm':
self.register_object(my_object.Round_peg_16mm_Object,
'peg',
xrange=[-0.1, -0.13],
yrange=[0.3, 0.33])
elif self.task_configs['peg'] == '12mm':
self.register_object(my_object.Round_peg_12mm_Object,
'peg',
xrange=[-0.1, -0.13],
yrange=[0.3, 0.33])
elif self.task_configs['peg'] == '9mm':
raise NotImplementedError
elif self.task_configs['peg'] == 'cylinder_16mm':
from robosuite.models.objects.primitive import CylinderObject
self.peg = CylinderObject('peg', size=(0.007, 0.025))
self.objects_of_interest.append(self.peg)
self.objectsName_of_interest.append('peg')
self.objectsXrange_of_interest.append([-0.1, -0.13])
self.objectsYrange_of_interest.append([0.3, 0.33])
# Create Sequential Sampler. The order is same as the order of register.
# Create individual samplers per object
self.placement_initializer = SequentialCompositeSampler(
name="ObjectSampler")
for obj_name, default_xy_range in zip(
self.objectsName_of_interest,
zip(self.objectsXrange_of_interest,
self.objectsYrange_of_interest)):
self.placement_initializer.append_sampler(
sampler=UniformRandomSampler(
name=f"{obj_name}Sampler",
x_range=default_xy_range[0],
y_range=default_xy_range[1],
rotation=None,
rotation_axis='z',
ensure_object_boundary_in_range=True,
ensure_valid_placement=True,
reference_pos=self.table_offset,
z_offset=0.01,
))
# Add objects to the sampler
for obj_to_put, obj_name in zip(self.objects_of_interest,
self.objectsName_of_interest):
self.placement_initializer.add_objects_to_sampler(
sampler_name=f"{obj_name}Sampler", mujoco_objects=obj_to_put)
if self.task_configs['board'] == 'hole':
self.plate = PlateWithHoleObject(name='plate')
plate_obj = self.plate.get_obj()
plate_obj.set("quat", "0 0 0 1")
plate_obj.set("pos", "0 0 {}".format(self.table_offset[2]))
self.objects_of_interest.append(self.plate)
# task includes arena, robot, and objects of interest
self.model = ManipulationTask(
mujoco_arena=mujoco_arena,
mujoco_robots=[robot.robot_model for robot in self.robots],
mujoco_objects=self.objects_of_interest,
)