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 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 _load_model(self):
super()._load_model()
self.mujoco_robot.set_base_xpos([0, 0, 0])
# load model for table top workspace
self.mujoco_arena = PegsArena(table_full_size=self.table_full_size,
table_friction=self.table_friction)
if self.use_indicator_object:
self.mujoco_arena.add_pos_indicator()
# The sawyer robot has a pedestal, we want to align it with the table
self.mujoco_arena.set_origin([.5, -0.15, 0])
# define mujoco objects
self.ob_inits = [SquareNutObject, RoundNutObject]
self.item_names = ["SquareNut", "RoundNut"]
self.item_names_org = list(self.item_names)
self.obj_to_use = (self.item_names[1] + "{}").format(0)
self.ngeoms = [5, 9]
lst = []
for i in range(len(self.ob_inits)):
ob = self.ob_inits[i]()
lst.append((str(self.item_names[i]) + "0", ob))
self.mujoco_objects = OrderedDict(lst)
self.n_objects = len(self.mujoco_objects)
# task includes arena, robot, and objects of interest
self.model = NutAssemblyTask(
self.mujoco_arena,
self.mujoco_robot,
self.mujoco_objects,
self.placement_initializer,
)
self.model.place_objects()
self.table_pos = string_to_array(self.model.table_body.get("pos"))
self.peg1_pos = string_to_array(
self.model.peg1_body.get("pos")) # square
self.peg2_pos = string_to_array(
self.model.peg2_body.get("pos")) # round
def _load_model(self):
super()._load_model()
self.mujoco_robot.set_base_xpos([0, 0, 0])
# load model for table top workspace
self.mujoco_arena = BinsArena(table_full_size=self.table_full_size,
table_friction=self.table_friction)
if self.use_indicator_object:
self.mujoco_arena.add_pos_indicator()
# The sawyer robot has a pedestal, we want to align it with the table
self.mujoco_arena.set_origin([.5, -0.3, 0])
self.ob_inits = [MilkObject, BreadObject, CerealObject, CanObject]
self.vis_inits = [
MilkVisualObject,
BreadVisualObject,
CerealVisualObject,
CanVisualObject,
]
self.item_names = ["Milk", "Bread", "Cereal", "Can"]
self.item_names_org = list(self.item_names)
self.obj_to_use = (self.item_names[0] + "{}").format(0)
lst = []
for j in range(len(self.vis_inits)):
lst.append((str(self.vis_inits[j]), self.vis_inits[j]()))
self.visual_objects = lst
lst = []
for i in range(len(self.ob_inits)):
ob = self.ob_inits[i]()
lst.append((str(self.item_names[i]) + "0", ob))
self.mujoco_objects = OrderedDict(lst)
self.n_objects = len(self.mujoco_objects)
# task includes arena, robot, and objects of interest
self.model = PickPlaceTask(
self.mujoco_arena,
self.mujoco_robot,
self.mujoco_objects,
self.visual_objects,
)
self.model.place_objects()
self.model.place_visual()
self.bin_pos = string_to_array(self.model.bin2_body.get("pos"))
self.bin_size = self.model.bin_size
def get_horizontal_radius(self):
horizontal_radius_site = self.worldbody.find(
"./body/site[@name='horizontal_radius_site']")
return string_to_array(horizontal_radius_site.get("pos"))[0]
def get_top_offset(self):
top_site = self.worldbody.find("./body/site[@name='top_site']")
return string_to_array(top_site.get("pos"))
def get_bottom_offset(self):
bottom_site = self.worldbody.find("./body/site[@name='bottom_site']")
return string_to_array(bottom_site.get("pos"))
def table_top_abs(self):
"""Returns the absolute position of table top"""
table_height = np.array([0, 0, self.table_full_size[2]])
return string_to_array(self.floor.get("pos")) + table_height
def table_top_abs(self):
"""
Returns the absolute position of table top.
"""
return string_to_array(self.table_body.get("pos"))
def bin_abs(self):
"""Returns the absolute position of table top"""
return string_to_array(self.bin1_body.get("pos"))