key_handler = KeyboardHandler(env=env, cam_body_id=cam_body_id) env.viewer.add_keypress_callback("any", key_handler.on_press) env.viewer.add_keyup_callback("any", key_handler.on_release) env.viewer.add_keyrepeat_callback("any", key_handler.on_press) # just spin to let user interact with glfw window spin_count = 0 while True: action = np.zeros(env.action_dim) obs, reward, done, _ = env.step(action) env.render() spin_count += 1 if spin_count % 500 == 0: # convert from world coordinates to file coordinates (xml subtree) camera_pos = np.array(env.sim.model.body_pos[cam_body_id]) camera_quat = T.convert_quat(env.sim.model.body_quat[cam_body_id], to='xyzw') world_camera_pose = T.make_pose(camera_pos, T.quat2mat(camera_quat)) file_camera_pose = world_in_file.dot(world_camera_pose) # TODO: Figure out why numba causes black screen of death (specifically, during mat2pose --> mat2quat call below) camera_pos, camera_quat = T.mat2pose(file_camera_pose) camera_quat = T.convert_quat(camera_quat, to='wxyz') print("\n\ncurrent camera tag you should copy") cam_tree.set("pos", "{} {} {}".format(camera_pos[0], camera_pos[1], camera_pos[2])) cam_tree.set("quat", "{} {} {} {}".format(camera_quat[0], camera_quat[1], camera_quat[2], camera_quat[3])) print(ET.tostring(cam_tree, encoding="utf8").decode("utf8"))
def _get_observation(self): """ Returns an OrderedDict containing observations [(name_string, np.array), ...]. Important keys: robot-state: contains robot-centric information. object-state: requires @self.use_object_obs to be True. contains object-centric information. image: requires @self.use_camera_obs to be True. contains a rendered frame from the simulation. depth: requires @self.use_camera_obs and @self.camera_depth to be True. contains a rendered depth map from the simulation """ di = super()._get_observation() if self.use_camera_obs: camera_obs = self.sim.render( camera_name=self.camera_name, width=self.camera_width, height=self.camera_height, depth=self.camera_depth, ) if self.camera_depth: di["image"], di["depth"] = camera_obs else: di["image"] = camera_obs # low-level object information if self.use_object_obs: # remember the keys to collect into object info object_state_keys = [] # for conversion to relative gripper frame gripper_pose = T.pose2mat((di["eef_pos"], di["eef_quat"])) world_pose_in_gripper = T.pose_inv(gripper_pose) for i in range(len(self.item_names_org)): if self.single_object_mode == 2 and self.object_id != i: # Skip adding to observations continue obj_str = str(self.item_names_org[i]) + "0" obj_pos = np.array( self.sim.data.body_xpos[self.obj_body_id[obj_str]]) obj_quat = T.convert_quat( self.sim.data.body_xquat[self.obj_body_id[obj_str]], to="xyzw") di["{}_pos".format(obj_str)] = obj_pos di["{}_quat".format(obj_str)] = obj_quat # get relative pose of object in gripper frame object_pose = T.pose2mat((obj_pos, obj_quat)) rel_pose = T.pose_in_A_to_pose_in_B(object_pose, world_pose_in_gripper) rel_pos, rel_quat = T.mat2pose(rel_pose) di["{}_to_eef_pos".format(obj_str)] = rel_pos di["{}_to_eef_quat".format(obj_str)] = rel_quat object_state_keys.append("{}_pos".format(obj_str)) object_state_keys.append("{}_quat".format(obj_str)) object_state_keys.append("{}_to_eef_pos".format(obj_str)) object_state_keys.append("{}_to_eef_quat".format(obj_str)) if self.single_object_mode == 1: # Zero out other objects observations for obj_str, obj_mjcf in self.mujoco_objects.items(): if obj_str == self.obj_to_use: continue else: di["{}_pos".format(obj_str)] *= 0.0 di["{}_quat".format(obj_str)] *= 0.0 di["{}_to_eef_pos".format(obj_str)] *= 0.0 di["{}_to_eef_quat".format(obj_str)] *= 0.0 di["object-state"] = np.concatenate( [di[k] for k in object_state_keys]) return di
def do_physical_simulation(_sim, _object_list, _action, _next_object_list, _viewer=None, test_horizon=20000, reset=True): state = _sim.get_state() if reset: for _obj in _object_list: if 'custom_table' not in _obj.name: pos_arr, quat_arr = T.mat2pose(_obj.pose) obj_qpos_addr = _sim.model.get_joint_qpos_addr(_obj.name) state.qpos[obj_qpos_addr[0]:obj_qpos_addr[0] + 3] = pos_arr state.qpos[obj_qpos_addr[0] + 3:obj_qpos_addr[0] + 7] = quat_arr[[3, 0, 1, 2]] if action['type'] == "pick": pass elif action['type'] == "place": held_obj_idx = get_held_object(_object_list) _obj = _next_object_list[held_obj_idx] pos_arr, quat_arr = T.mat2pose(deepcopy(_obj.pose)) pos_arr[2] += 0.02 obj_qpos_addr = _sim.model.get_joint_qpos_addr(_obj.name) state.qpos[obj_qpos_addr[0]:obj_qpos_addr[0] + 3] = pos_arr state.qpos[obj_qpos_addr[0] + 3:obj_qpos_addr[0] + 7] = quat_arr[[3, 0, 1, 2]] _sim.set_state(state) for _ in range(test_horizon): _sim.step() if _viewer is not None: _viewer.render() state = _sim.get_state() mis_posed_object_list = [] sim_object_list = [] for _obj in _next_object_list: sim_obj = deepcopy(_obj) if 'custom_table' not in _obj.name: pos_arr, quat_arr = T.mat2pose(_obj.pose) obj_qpos_addr = _sim.model.get_joint_qpos_addr(_obj.name) sim_pos_arr = state.qpos[obj_qpos_addr[0]:obj_qpos_addr[0] + 3] sim_quat_arr = state.qpos[obj_qpos_addr[0] + 3:obj_qpos_addr[0] + 7] sim_obj.pose = T.pose2mat( [sim_pos_arr, sim_quat_arr[[1, 2, 3, 0]]]) dist_diff = np.sqrt(np.sum((pos_arr - sim_pos_arr)**2)) ang_diff = np.arccos( np.maximum( np.minimum( 2. * np.sum(quat_arr * sim_quat_arr[[1, 2, 3, 0]])**2 - 1., 1.), -1.)) if dist_diff > 0.05 or ang_diff > np.pi / 20.: mis_posed_object_list.append(_obj) sim_object_list.append(sim_obj) return mis_posed_object_list, sim_object_list
color=obj.color) if goal_obj is not None: ax = visualize_trimesh(meshes[goal_obj.mesh_idx], goal_obj.pose, ax=ax, color=goal_obj.color) plt.show() arena_model = MujocoXML( xml_path_completion("arenas/empty_arena.xml")) for obj in initial_object_list: if 'custom_table' in obj.name: table_model = MujocoXML( xml_path_completion("objects/custom_table.xml")) table_pos_arr, table_quat_arr = T.mat2pose(obj.pose) table_body = table_model.worldbody.find( "./body[@name='custom_table']") table_body.set("pos", array_to_string(table_pos_arr)) table_body.set( "quat", array_to_string(table_quat_arr[[3, 0, 1, 2]])) arena_model.merge(table_model) else: if 'arch_box' in obj.name: obj_xml_path = "objects/arch_box.xml" elif 'rect_box' in obj.name: obj_xml_path = "objects/rect_box.xml" elif 'square_box' in obj.name: obj_xml_path = "objects/square_box.xml" elif 'half_cylinder_box' in obj.name: obj_xml_path = "objects/half_cylinder_box.xml"
def _get_observation(self): """ Returns an OrderedDict containing observations [(name_string, np.array), ...]. Important keys: `'robot-state'`: contains robot-centric information. `'object-state'`: requires @self.use_object_obs to be True. Contains object-centric information. `'image'`: requires @self.use_camera_obs to be True. Contains a rendered frame from the simulation. `'depth'`: requires @self.use_camera_obs and @self.camera_depth to be True. Contains a rendered depth map from the simulation Returns: OrderedDict: Observations from the environment """ di = super()._get_observation() # low-level object information if self.use_object_obs: # Get robot prefix pr = self.robots[0].robot_model.naming_prefix # remember the keys to collect into object info object_state_keys = [] # for conversion to relative gripper frame gripper_pose = T.pose2mat( (di[pr + "eef_pos"], di[pr + "eef_quat"])) world_pose_in_gripper = T.pose_inv(gripper_pose) for i in range(len(self.item_names_org)): if self.single_object_mode == 2 and self.nut_id != i: # skip observations continue obj_str = str(self.item_names_org[i]) + "0" obj_pos = np.array( self.sim.data.body_xpos[self.obj_body_id[obj_str]]) obj_quat = T.convert_quat( self.sim.data.body_xquat[self.obj_body_id[obj_str]], to="xyzw") di["{}_pos".format(obj_str)] = obj_pos di["{}_quat".format(obj_str)] = obj_quat object_pose = T.pose2mat((obj_pos, obj_quat)) rel_pose = T.pose_in_A_to_pose_in_B(object_pose, world_pose_in_gripper) rel_pos, rel_quat = T.mat2pose(rel_pose) di["{}_to_{}eef_pos".format(obj_str, pr)] = rel_pos di["{}_to_{}eef_quat".format(obj_str, pr)] = rel_quat object_state_keys.append("{}_pos".format(obj_str)) object_state_keys.append("{}_quat".format(obj_str)) object_state_keys.append("{}_to_{}eef_pos".format(obj_str, pr)) object_state_keys.append("{}_to_{}eef_quat".format( obj_str, pr)) if self.single_object_mode == 1: # zero out other objs for obj_str, obj_mjcf in self.mujoco_objects.items(): if obj_str == self.obj_to_use: continue else: di["{}_pos".format(obj_str)] *= 0.0 di["{}_quat".format(obj_str)] *= 0.0 di["{}_to_{}eef_pos".format(obj_str, pr)] *= 0.0 di["{}_to_{}eef_quat".format(obj_str, pr)] *= 0.0 di["object-state"] = np.concatenate( [di[k] for k in object_state_keys]) return di
mcts = Tree(initial_object_list, np.sum(n_obj_per_mesh_types) * 2, coll_mngr, meshes, contact_points, contact_faces, rotation_types, _goal_obj=goal_obj, _physcial_constraint_checker=None) for _ in range(10): mcts.exploration(0) print("Planning") optimized_path = mcts.get_best_path(0) final_object_list = mcts.Tree.nodes[optimized_path[-1]]['state'] print("Do dynamic simulation") if len(configuration_list) == 0: arena_model = MujocoXML(xml_path_completion("arenas/empty_arena.xml")) for obj in initial_object_list: if 'custom_table' in obj.name: table_model = MujocoXML(xml_path_completion("objects/custom_table.xml")) table_pos_arr, table_quat_arr = T.mat2pose(obj.pose) table_body = table_model.worldbody.find("./body[@name='custom_table']") table_body.set("pos", array_to_string(table_pos_arr)) table_body.set("quat", array_to_string(table_quat_arr[[3, 0, 1, 2]])) arena_model.merge(table_model) else: if 'arch_box' in obj.name: obj_xml_path = "objects/arch_box.xml" elif 'rect_box' in obj.name: obj_xml_path = "objects/rect_box.xml" elif 'square_box' in obj.name: obj_xml_path = "objects/square_box.xml" elif 'half_cylinder_box' in obj.name: obj_xml_path = "objects/half_cylinder_box.xml" elif 'triangle_box' in obj.name: obj_xml_path = "objects/triangle_box.xml"