def _env_setup(self, initial_qpos): for name, value in initial_qpos.items(): self.sim.data.set_joint_qpos(name, value) utils.reset_mocap_welds(self.sim) self.sim.forward() # Move end effector into position. target = self.sim.data.get_site_xpos('needle_tip').copy() self.rotation = self.sim.data.get_body_xquat('needle_tip').copy() self.sim.data.set_mocap_pos('kuka_mocap', target) self.sim.data.set_mocap_quat('kuka_mocap', self.rotation) # sim forward for _ in range(10): self.sim.step() sites_offset = (self.sim.data.site_xpos - self.sim.model.site_pos).copy()[3] self.target_center = self.sim.data.get_site_xpos('target_center') self.init_center = self.sim.data.get_site_xpos('init_center') self.initial_needle_xpos = self.sim.data.get_site_xpos( 'needle_tip').copy() self.height_offset = self.sim.data.get_site_xpos('object0')[2] site_id = self.sim.model.site_name2id('init_1') self.sim.model.site_pos[site_id] = self.init_center + [ self.obj_range, self.obj_range, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('init_2') self.sim.model.site_pos[site_id] = self.init_center + [ self.obj_range, -self.obj_range, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('init_3') self.sim.model.site_pos[site_id] = self.init_center + [ -self.obj_range, self.obj_range, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('init_4') self.sim.model.site_pos[site_id] = self.init_center + [ -self.obj_range, -self.obj_range, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('mark1') self.sim.model.site_pos[site_id] = self.target_center + [ self.target_range_x, self.target_range_y, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('mark2') self.sim.model.site_pos[site_id] = self.target_center + [ -self.target_range_x, self.target_range_y, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('mark3') self.sim.model.site_pos[site_id] = self.target_center + [ self.target_range_x, -self.target_range_y, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('mark4') self.sim.model.site_pos[site_id] = self.target_center + [ -self.target_range_x, -self.target_range_y, 0.0 ] - sites_offset self.sim.step() if self.has_object: self.height_offset = self.sim.data.get_site_xpos('object0')[2]
def _env_setup(self, initial_qpos): for name, value in initial_qpos.items(): self.sim.data.set_joint_qpos(name, value) utils.reset_mocap_welds(self.sim) self.sim.forward() # Move end effector into position. target = self.sim.data.get_site_xpos('gripper_tip').copy() rotation = self.sim.data.get_body_xquat('gripper_tip').copy() self.sim.data.set_mocap_pos('kuka_mocap', target) self.sim.data.set_mocap_quat('kuka_mocap', rotation) for _ in range(10): self.sim.step() # initial markers (index 3 is arbitrary) self.target_1 = self.sim.data.get_site_xpos('target_1') self.target_2 = self.sim.data.get_site_xpos('target_2') self.target_3 = self.sim.data.get_site_xpos('target_3') self.target_4 = self.sim.data.get_site_xpos('target_4') self.target_5 = self.sim.data.get_site_xpos('target_5') self.target_6 = self.sim.data.get_site_xpos('target_6') self.target_7 = self.sim.data.get_site_xpos('target_7') self.target_8 = self.sim.data.get_site_xpos('target_8') self.targets = [ self.target_1, self.target_2, self.target_3, self.target_4, self.target_5, self.target_6, self.target_7, self.target_8 ] self.init_center = self.sim.data.get_site_xpos('init_center') site_id = self.sim.model.site_name2id('init_center') sites_offset = (self.sim.data.site_xpos[site_id] - self.sim.model.site_pos[site_id]).copy() self.initial_gripper_xpos = self.sim.data.get_site_xpos( 'gripper_tip').copy() site_id = self.sim.model.site_name2id('init_1') self.sim.model.site_pos[site_id] = self.init_center + [ self.obj_range, self.obj_range, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('init_2') self.sim.model.site_pos[site_id] = self.init_center + [ self.obj_range, -self.obj_range, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('init_3') self.sim.model.site_pos[site_id] = self.init_center + [ -self.obj_range, self.obj_range, 0.0 ] - sites_offset site_id = self.sim.model.site_name2id('init_4') self.sim.model.site_pos[site_id] = self.init_center + [ -self.obj_range, -self.obj_range, 0.0 ] - sites_offset self.sim.step() if self.has_object: self.height_offset = self.sim.data.get_site_xpos('object0')[2]
def _env_setup(self, initial_qpos): for name, value in initial_qpos.items(): self.sim.data.set_joint_qpos(name, value) utils.reset_mocap_welds(self.sim) self.sim.forward() target = self.sim.data.get_site_xpos('needle_tip').copy() self.rotation = self.sim.data.get_body_xquat('needle_tip').copy() self.sim.data.set_mocap_pos('kuka_mocap', target) self.sim.data.set_mocap_quat('kuka_mocap', self.rotation) for _ in range(50): self.sim.step() self.initial_needle_xpos = self.sim.data.get_site_xpos( 'needle_tip').copy() if self.has_object: self.height_offset = self.sim.data.get_site_xpos('object0')[2]