# -*- coding: utf-8 -*- import json import time import numpy as np from scipy.spatial.transform import Rotation as R, Slerp from data_gen_dependencies.robot import IsaacSimRpcRobot class BaseAgent(object): def __init__(self, robot: IsaacSimRpcRobot) -> None: self.robot = robot self.reset() def reset(self): self.robot.reset() # self.base_camera = "/World/Top_Camera" # self.robot.base_camera = "/World/Top_Camera" # self.data_keys = { # "camera": { # "camera_prim_list": [self.base_camera], # "render_depth": True, # "render_semantic": True, # }, # "pose": [self.base_camera], # "joint_position": True, # "gripper": True, # } # self.get_observation() def add_object(self, object_info: dict): name = object_info["object_id"] usd_path = object_info["model_path"] position = np.array(object_info["position"]) quaternion = np.array(object_info["quaternion"]) if "scale" not in object_info: object_info["scale"] = 1.0 size = object_info.get('size', np.array([])) if isinstance(size, list): size = np.array(size) if size.shape[0] == 0: size = np.array(100) if np.max(size) > 10: object_info["scale"] = 0.001 add_particle = False particle_position = [0, 0, 0] particle_scale = [1, 1, 1] if "add_particle" in object_info: add_particle = True particle_position = object_info["add_particle"]["position"] particle_scale = object_info["add_particle"]["scale"] scale = np.array([object_info["scale"]] * 3) material = "general" if "material" not in object_info else object_info["material"] mass = 0.01 if "mass" not in object_info else object_info["mass"] # if "materialOptions" in object_info: # if "transparent" in object_info["materialOptions"]: # material="Glass" self.robot.client.add_object( usd_path=usd_path, prim_path="/World/Objects/%s" % name, label_name=name, target_position=position, target_quaternion=quaternion, target_scale=scale, material=material, color=[1, 1, 1], mass=mass, add_particle=add_particle, particle_position=particle_position, particle_scale=particle_scale ) def generate_layout(self, task_file, init=True): with open(task_file, "r") as f: task_info = json.load(f) # init_position = task_info["init_position"] if "init_position" in task_info else [0] * 39 # self.robot.client.set_joint_positions(init_position) if init: for object_info in task_info["objects"]: if "box" not in object_info["name"]: continue self.add_object(object_info) time.sleep(1) for object_info in task_info["objects"]: if "obj" not in object_info["name"]: continue self.add_object(object_info) time.sleep(1) self.task_info = task_info self.robot.target_object = task_info["target_object"] self.arm = task_info["arm"] return task_info def execute(self, commands): for command in commands: type = command["action"] content = command["content"] if type == "move" or type == "rotate": if self.robot.move(content) == False: return False else: self.last_pose = content["matrix"] elif type == "open_gripper": id = "left" if "gripper" not in content else content["gripper"] width = 0.1 if "width" not in content else content["width"] self.robot.open_gripper(id=id, width=width) elif type == "close_gripper": id = "left" if "gripper" not in content else content["gripper"] force = 50 if "force" not in content else content["force"] self.robot.close_gripper(id=id, force=force) if "attach_obj" in content: self.robot.client.AttachObj(prim_paths=[content["attach_obj"]]) else: Exception("Not implemented.") return True def start_recording(self, task_name, camera_prim_list): print(camera_prim_list) self.robot.client.start_recording( task_name=task_name, fps=30, data_keys={ "camera": { "camera_prim_list": camera_prim_list, "render_depth": False, "render_semantic": False, }, "pose": ["/World/Raise_A2/gripper_center"], "joint_position": True, "gripper": True, }, ) def stop_recording(self, success): self.robot.client.stop_recording() self.robot.client.SendTaskStatus(success) def grasp( self, target_gripper_pose, gripper_id=None, use_pre_grasp=False, use_pick_up=False, grasp_width=0.1, ): gripper_id = "left" if gripper_id is None else gripper_id pick_up_pose = np.copy(target_gripper_pose) pick_up_pose[2, 3] = pick_up_pose[2, 3] + 0.1 # 抓到物体后，先往上方抬10cm commands = [] commands.append( { "action": "open_gripper", "content": { "gripper": gripper_id, "width": grasp_width, }, } ) if use_pre_grasp: pre_pose = np.array([ [1.0, 0, 0, 0], [0, 1.0, 0, 0], [0, 0, 1.0, -0.05], [0, 0, 0, 1] ]) pre_grasp_pose = target_gripper_pose @ pre_pose commands.append( { "action": "move", "content": { "matrix": pre_grasp_pose, "type": "matrix", "comment": "pre_grasp", "trajectory_type": "ObsAvoid", }, } ) grasp_trajectory_type = "Simple" if use_pre_grasp else "ObsAvoid" commands.append( { "action": "move", "content": { "matrix": target_gripper_pose, "type": "matrix", "comment": "grasp", "trajectory_type": grasp_trajectory_type, }, } ) commands.append( { "action": "close_gripper", "content": { "gripper": gripper_id, "force": 50, }, } ) if use_pick_up: commands.append( { "action": "move", "content": { "matrix": pick_up_pose, "type": "matrix", "comment": "pick_up", "trajectory_type": "Simple", }, } ) return commands def get_observation(self): observation_raw = self.robot.get_observation(self.data_keys) self.observation = observation_raw["camera"][self.robot.base_camera] def place( self, target_gripper_pose, current_gripper_pose=None, gripper2part=None, gripper_id=None, ): gripper_id = "left" if gripper_id is None else gripper_id pre_place_pose = np.copy(target_gripper_pose) pre_place_pose[2, 3] += 0.05 commands = [ { "action": "move", "content": { "matrix": pre_place_pose, "type": "matrix", "comment": "pre_place", "trajectory_type": "ObsAvoid", }, }, { "action": "move", "content": { "matrix": target_gripper_pose, "type": "matrix", "comment": "place", "trajectory_type": "Simple", }, }, { "action": "open_gripper", "content": { "gripper": gripper_id, }, }, ] return commands def pour( self, target_gripper_pose, current_gripper_pose=None, gripper2part=None, gripper_id=None, ): def interpolate_rotation_matrices(rot_matrix1, rot_matrix2, num_interpolations): rot1 = R.from_matrix(rot_matrix1) rot2 = R.from_matrix(rot_matrix2) quat1 = rot1.as_quat() quat2 = rot2.as_quat() times = [0, 1] slerp = Slerp(times, R.from_quat([quat1, quat2])) interp_times = np.linspace(0, 1, num_interpolations) interp_rots = slerp(interp_times) interp_matrices = interp_rots.as_matrix() return interp_matrices target_part_pose = target_gripper_pose @ np.linalg.inv(gripper2part) current_part_pose = current_gripper_pose @ np.linalg.inv(gripper2part) commands = [] rotations = interpolate_rotation_matrices(current_part_pose[:3, :3], target_part_pose[:3, :3], 5) for i, rotation in enumerate(rotations): target_part_pose_step = np.copy(target_part_pose) target_part_pose_step[:3, :3] = rotation target_gripper_pose_step = target_part_pose_step @ gripper2part commands.append( { "action": "move", "content": { "matrix": target_gripper_pose_step, "type": "matrix", "comment": "pour_sub_rotate_%d" % i, "trajectory_type": "Simple", }, } ) return commands