update readme

.gitignore

@@ -2,4 +2,5 @@ __pycache__/
 workspace/
 logs/
 cache/
-.idea/
+.idea/
+*.txt

README.md

@@ -0,0 +1,112 @@
+# Isaac Sim Agent
+
+Isaac Sim Agent 是一个基于 NVIDIA Isaac Sim 的机器人任务生成和数据采集系统。该系统支持任务模板划分、任务实例生成、数据采集和回放功能。
+
+## 系统架构
+
+系统包含以下主要组件：
+- **任务模板划分器** (TaskTemplatesDivider): 将大型任务模板划分为多个子任务
+- **任务生成器** (TaskGenerator): 从任务模板生成具体的任务实例
+- **数据生成器** (DataGenerator): 在 Isaac Sim 中执行任务并采集数据
+- **数据回放器** (DataReplayer)（暂未接入）: 回放已采集的数据
+- **回放初始化器** (ReplayInitializer)（暂未接入）: 为数据回放准备环境
+
+## 1 安装依赖
+
+安装以下依赖：
+```bash
+pip install pin -i https://pypi.tuna.tsinghua.edu.cn/simple
+pip install numpy==1.26 -i https://pypi.tuna.tsinghua.edu.cn/simple
+pip install scikit-learn -i https://pypi.tuna.tsinghua.edu.cn/simple
+pip install grasp_nms -i https://pypi.tuna.tsinghua.edu.cn/simple
+pip install open3d -i https://pypi.tuna.tsinghua.edu.cn/simple
+pip install git+https://git.hofee.cn/hofee/pyboot.git 
+```
+
+执行 `pb scan` 检查依赖是否完整，若缺少某一依赖，安装对应依赖。
+
+## 2 配置说明
+
+### 2.1 任务模板划分配置 (divide_task_config.yaml)
+```yaml
+runner:
+  workspace:
+    name: "divide_task"
+    root_dir: "workspace"
+  divide:
+    input_task_templates_root_dir: "/path/to/input/task/templates"
+    output_task_templates_dir: ""  # 为空则默认在workspace/divide_task/task_templates
+    output_template_targets_dir: ""  # 为空则默认在workspace/divide_task/template_targets
+    divide_num: 100  # 每个进程划分任务数量
+    total_nums: 10000  # 总任务模板数量
+```
+
+### 2.2 任务生成配置 (generate_task_config.yaml)
+```yaml
+runner:
+  workspace:
+    name: "generate_task"
+    root_dir: "workspace"
+  generate:
+    input_data_root: "/path/to/input/data"
+    input_target_task_templates_path: "workspace/divide_task/template_targets/task_template_target_0.json"
+    output_task_root_dir: ""  # 为空则默认在workspace/generate_task/task_root_dir
+    output_generate_result_path: ""  # 为空则默认在workspace/generate_task/generate_result_path.json
+```
+
+### 2.3 数据生成配置 (generate_data_config.yaml)
+```yaml
+runner:
+  workspace:
+    name: "generate_data"
+    root_dir: "workspace"
+  generate:
+    input_data_root: "/path/to/input/data"
+    input_target_task_path: "workspace/generate_task/generate_result_path.json"
+    server_url: "localhost:50051"  # Isaac Sim 服务器地址
+```
+
+## 3 使用流程
+
+### 3.1 划分任务模板
+在 `configs/divide_task_config.yaml` 里配置任务模板输入目录、整理好的任务模板输出目录（为空则默认在workspace/divide_task/task_templates）、用于多机生产的划分后的Task Template列表（为空则默认在workspace/divide_task/template_targets）
+
+```bash
+pb run divide_task
+```
+
+### 3.2 生成任务实例
+从指定任务模版生成任务实例：
+
+```bash
+pb run generate_task
+```
+
+### 3.3 生成数据
+从指定任务实例表里依次运行任务实例（目前只支持运行同一个任务模板下的多个任务实例）：
+
+```bash
+pb run generate_data
+```
+
+### 3.4 数据回放（暂未接入）
+回放已生成的数据：
+
+```bash
+pb run replay
+```
+
+## 4 工作空间结构
+
+```
+workspace/
+├── divide_task/          # 任务模板划分输出
+│   ├── task_templates/   # 整理后的任务模板
+│   └── template_targets/ # 划分后的任务模板目标
+├── generate_task/        # 任务生成输出
+│   ├── task_root_dir/    # 生成的任务实例
+│   └── generate_result_path.json  # 生成结果路径
+└── generate_data/        # 数据生成输出
+    └── [生成的数据文件]
+```
+

app.py

@@ -1,7 +1,7 @@
 from pyboot.application import PybootApplication
 from runners.task_generator import TaskGenerator
 from runners.data_generator import DataGenerator
-from runners.data_recorder import DataRecorder
+from runners.data_replayer import DataReplayer
 from runners.task_templates_divider import TaskTemplatesDivider
 
 @PybootApplication("main")
@@ -10,7 +10,7 @@ class MainApp:
     def start():
         TaskGenerator(config_path="configs/generate_task_config.yaml").run()
         DataGenerator(config_path="configs/generate_data_config.yaml").run()
-        DataRecorder(config_path="configs/record_data_config.yaml").run()
+        DataReplayer(config_path="configs/record_data_config.yaml").run()
 
 @PybootApplication("divide_task")
 class DivideTaskApp:
@@ -30,8 +30,8 @@ class GenerateDataApp:
     def start():
         DataGenerator(config_path="configs/generate_data_config.yaml").run()
 
-@PybootApplication("record_data")
-class RecordDataApp:
+@PybootApplication("replay")
+class ReplayApp:
     @staticmethod
     def start():
-        DataRecorder(config_path="configs/record_data_config.yaml").run()
+        DataReplayer(config_path="configs/replay_config.yaml").run()

configs/generate_data_config.yaml

@@ -4,4 +4,5 @@ runner:
     root_dir: "workspace"
   generate:
     input_data_root: "/home/ubuntu/Projects/docker_isaac_sim/input/data"
-    input_target_task_templates_path: "workspace/divide_task/template_targets/task_template_target_0.json"
+    input_target_task_path: "workspace/generate_task/generate_result_path.json"
+    server_url: "localhost:50051"

data_gen_dependencies/action/__init__.py

@@ -4,7 +4,11 @@ from .insert import InsertStage, HitStage
 from .slide import SlideStage
 from .pour import PourStage
 from .pull import PullStage
-from .push import PushStage, ClickStage
+from .push import PushStage
+from .pull_revolute import PullRevoluteStage
+from .push_revolute import PushRevoluteStage
+from .press import PressPrismaticStage
+from .twist import TwistStage
 
 ACTION_STAGE = {
     "grasp": GraspStage,
@@ -19,9 +23,16 @@ ACTION_STAGE = {
     "hit": NotImplemented,
     "pour": PourStage,
     "push": PushStage,
-    'click': ClickStage,
-    'touch': ClickStage,
-    "pull": PullStage
+    'click': PressPrismaticStage,
+    'touch': PressPrismaticStage,
+    "pull": PullStage,
+    # ---- Articulate ----
+    "pull_revolute": PullRevoluteStage,
+    "push_revolute": PushRevoluteStage,
+    "pull_prismatic":PullStage,
+    "push_prismatic":PushStage,
+    "twist": TwistStage,
+    'press_prismatic': PressPrismaticStage,
 }
 
 def build_stage(action):

data_gen_dependencies/action/base.py

@@ -27,7 +27,6 @@ def simple_check_completion(goal, objects, last_statement=None, pos_threshold=0.
 
 def solve_target_gripper_pose(stage, objects):
     active_obj_ID, passive_obj_ID, target_pose_canonical, gripper_action, transform_world, motion_type = stage
-    
     anchor_pose = objects[passive_obj_ID].obj_pose
     
     
@@ -83,7 +82,7 @@ class StageTemplate:
             delta_pose = gripper_pose_canonical
             gripper_pose = objects['gripper'].obj_pose
             target_gripper_pose = gripper_pose @ delta_pose
-            motion_type = 'Straight'
+            motion_type = 'Simple'
         else:
 
             if gripper_pose_canonical is None:

data_gen_dependencies/action/grasp.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, simple_check_comple
 
 
 class PickStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element, passive_element, grasp_pose, extra_params=None, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element, passive_element, grasp_pose, extra_params=None, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
         self.grasp_pose_canonical = grasp_pose
         self.use_pre_grasp = False
@@ -104,6 +104,7 @@ class GraspStage(PickStage):
             
             self.sub_stages.append([pre_grasp_pose, None, np.eye(4), motion_type])
         # sub-stage-1   moveTo grasp pose
+
         self.sub_stages.append([grasp_pose, 'close', np.eye(4), motion_type])
 
         

data_gen_dependencies/action/insert.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, simple_check_comple
 
 
 class InsertStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None, extra_params={}, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None, extra_params={}, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
 
 

data_gen_dependencies/action/place.py

@@ -4,27 +4,24 @@ from data_gen_dependencies.action.base import StageTemplate, simple_check_comple
 
 
 class PlaceStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, extra_params={}, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, extra_params={}, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
         self.pre_transform_up =0.12
         self.place_transform_up = np.array([0, 0, 0.01])
         self.extra_params = {} if extra_params is None else extra_params
         self.use_pre_place = self.extra_params.get('use_pre_place', True)
         self.generate_substage(target_pose)
-        
-        
- 
-        
+
     def generate_substage(self, target_pose):
         target_pose_canonical = target_pose
         gripper_cmd = self.extra_params.get('gripper_state', 'open')
         pre_place_direction = self.extra_params.get('pre_place_direction', 'z')
-        print(gripper_cmd)
+        #print(gripper_cmd)
         num_against = self.extra_params.get('against', 0)
         if self.use_pre_place:
             # moveTo pre-place position
             transform_up = np.eye(4)
-            print(transform_up)
+            #print(transform_up)
             if pre_place_direction == "x":
                 transform_up[0,3] = self.pre_transform_up
                 transform_up[2,3] = 0.02

data_gen_dependencies/action/pour.py

@@ -9,7 +9,7 @@ import numpy as np
 
 
 class PourStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, current_pose=None, obj2part=None, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, current_pose=None, obj2part=None, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
         
         self.pre_pour_height = 0.25

data_gen_dependencies/action/press.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+from data_gen_dependencies.action.base import StageTemplate
+
+class PressPrismaticStage(StageTemplate):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None, objects=[], **kwargs):
+        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
+        self.pre_distance = -0.03
+        
+        self.passive_obj = objects[passive_obj_id]
+        self.active_obj = objects[active_obj_id]
+        self.joint_position_threshold = passive_element.get('joint_position_threshold', 0.7)
+        self.joint_position_threshold = self.joint_position_threshold if self.joint_position_threshold is not None else 0.7
+        self.correspond_joint_id = passive_element.get('correspond_joint_id', None)
+        correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id]
+        self.joint_lower_limit = correspond_joint_info["lower_bound"]
+        self.joint_upper_limit = correspond_joint_info["upper_bound"]
+        self.move_distance = -0.005 + abs(self.joint_lower_limit - self.joint_upper_limit) * self.joint_position_threshold
+        self.joint_axis = correspond_joint_info["joint_axis"]
+        self.joint_type = correspond_joint_info["joint_type"]
+        assert self.joint_type == 'prismatic', "joint_type must be prismatic for press action"
+        correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id]
+        self.generate_substage(target_pose, vector_direction)
+
+        
+    def generate_substage(self, target_pose, vector_direction):
+        vector_direction = vector_direction / np.linalg.norm(vector_direction)
+        pre_pose = target_pose.copy()
+        pre_pose[:3,3] += vector_direction * self.pre_distance
+        self.sub_stages.append([pre_pose, 'close', np.eye(4), 'AvoidObs'])
+        
+        move_pose = target_pose.copy()
+        move_pose[:3,3] += vector_direction * self.move_distance
+        self.sub_stages.append([move_pose, None, np.eye(4), 'Simple'])
+        free_delta_pose = np.eye(4)
+        free_delta_pose[2,3] = -0.03
+        self.sub_stages.append([free_delta_pose, "open", np.eye(4), 'local_gripper'])

data_gen_dependencies/action/pull.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate
 
 
 class PullStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
         self.pull_distance = passive_element['distance']
         assert self.pull_distance > 0
@@ -32,8 +32,6 @@ class PullStage(StageTemplate):
         free_delta_pose[2,3] = -0.03
         self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper'])
 
-        
-
 
     def check_completion(self, objects):
         if self.__len__()==0:

data_gen_dependencies/action/pull_revolute.py

@@ -0,0 +1,145 @@
+import os
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+from data_gen_dependencies.action.base import StageTemplate
+
+
+class PullRevoluteStage(StageTemplate):
+    DELTA_DISTANCE = 0.003  # meter
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects=[], **kwargs):
+        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
+        #import ipdb;ipdb.set_trace()
+        self.objects = objects
+        self.passive_obj = objects[passive_obj_id]
+        self.active_obj = objects[active_obj_id]
+        main_part_name= passive_element["part_id"]
+        self.obj_base = os.path.dirname(passive_obj_id)
+        main_part_id = os.path.join(self.obj_base, main_part_name)
+        if main_part_id not in objects:
+            raise ValueError(f"Cannot find main part {main_part_name} in objects")
+        self.root_id = os.path.join(self.obj_base, "root")
+        self.main_part_obj = objects[main_part_id]
+        self.root_obj = objects[self.root_id]
+        self.revolute_radius = np.linalg.norm(self.passive_obj.obj_pose[:3,3] - self.main_part_obj.obj_pose[:3,3])
+        self.revolute_joint_position = self.main_part_obj.obj_pose[:3,3]
+        self.joint_position_threshold = passive_element.get('joint_position_threshold', 0.7)
+        self.correspond_joint_id = passive_element.get('correspond_joint_id', None)
+        correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id]
+        
+        self.joint_lower_limit = correspond_joint_info["lower_bound"]
+        self.joint_upper_limit = correspond_joint_info["upper_bound"]
+        self.joint_axis = correspond_joint_info["joint_axis"]
+        self.joint_type = correspond_joint_info["joint_type"]
+        assert self.joint_type == 'revolute', "joint_type must be revolute for pull_revolute action"
+        if self.joint_position_threshold is None:
+            self.joint_position_threshold = 0.7
+        assert self.joint_position_threshold >= 0 and self.joint_position_threshold <= 1 
+
+        self.joint_direction = passive_element.get('joint_direction', 1)
+        assert self.joint_direction in [-1, 1]
+
+        self.joint_velocity_threshold = passive_element.get('joint_velocity_threshold', 999)
+        vector_direction = passive_element['direction']
+
+        self.extra_params = {} if extra_params is None else extra_params
+        self.generate_substage(target_pose, vector_direction)
+
+    def generate_substage(self, gripper2obj, vector_direction):
+        vector_direction = vector_direction / np.linalg.norm(vector_direction)
+        total_angle = self.joint_position_threshold * (self.joint_upper_limit - self.joint_lower_limit)
+        arc_length = abs(self.revolute_radius * total_angle)
+        step = int(arc_length / self.DELTA_DISTANCE)
+        from pyboot.utils.log import Log
+        Log.debug(f'PullRevoluteStage: total_angle={total_angle}, arc_length={arc_length}, step={step}')
+        trajectory = self.generate_trajectory(gripper2obj, vector_direction)
+        self.save_visualized_gripper_trajectory(trajectory, "pull_revolute_traj.txt")
+        
+        self.sub_stages.append([trajectory, None, np.eye(4), 'Trajectory'])
+        self.sub_stages.append([None, "open", np.eye(4), 'Simple'])
+
+        free_delta_pose = np.eye(4)
+        free_delta_pose[2,3] = -0.03
+        self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper'])
+
+
+    def generate_trajectory(self, gripper2passive_obj, init_vector_direction):
+        trajectory = []
+        total_angle = self.joint_position_threshold * (self.joint_upper_limit - self.joint_lower_limit)
+        total_angle *= self.joint_direction 
+
+        arc_length = abs(self.revolute_radius * total_angle)
+        num_steps = int(np.ceil(arc_length / self.DELTA_DISTANCE))
+        if num_steps < 1:
+            num_steps = 1
+
+        delta_theta = total_angle / num_steps
+        rotation_axis = np.array(self.joint_axis, dtype=float)
+        rotation_axis /= np.linalg.norm(rotation_axis) 
+
+        passive_obj_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.passive_obj.obj_pose
+        gripper_2_root = passive_obj_2_root @ gripper2passive_obj
+        pos0 = gripper_2_root[:3, 3]
+        R0 = gripper_2_root[:3, :3]
+
+        joint_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.main_part_obj.obj_pose
+
+        pivot = joint_2_root[:3, 3]
+
+        for i in range(1, num_steps + 1):
+            angle = i * delta_theta
+            R_delta = R.from_rotvec(rotation_axis * angle).as_matrix()
+            pos_rot = R_delta @ (pos0 - pivot) + pivot
+            R_new = R_delta @ R0
+
+            waypoint_2_root = np.eye(4)
+            waypoint_2_root[:3, :3] = R_new
+            waypoint_2_root[:3, 3] = pos_rot
+            waypoint_2_obj = np.linalg.inv(passive_obj_2_root) @ waypoint_2_root
+            trajectory.append(waypoint_2_obj)
+        return np.asarray(trajectory)
+
+    def save_visualized_gripper_trajectory(self, gripper_trajectory, save_path):
+        #import ipdb; ipdb.set_trace()
+        N = len(gripper_trajectory)
+        gripper_pts = np.zeros((1, 60, 3))
+        for i in range(10):
+            gripper_pts[0, i] = np.array([0, 0.01, 0.002*i])
+            gripper_pts[0, i+10] = np.array([0, 0.01, 0.002*i + 0.001])
+            gripper_pts[0, i+20] = np.array([0, -0.01, 0.002*i])
+            gripper_pts[0, i+30] = np.array([0, -0.01, 0.002*i + 0.001])
+            gripper_pts[0, i+40] = np.array([0, -0.01 + 0.002 * i, 0])
+        for i in range(5):
+            gripper_pts[0, i+50] = np.array([0, 0, -0.002*i])
+            gripper_pts[0, i+55] = np.array([0, 0, -0.002*i + 0.001])
+        gripper_pts -= np.array([0, 0, 0.002*10 + 0.001])
+        gripper_pts = gripper_pts.repeat(N, axis=0)
+        gripper_trajectory = gripper_trajectory.reshape(N, 1, 4, 4)
+        gripper_trajectory = gripper_trajectory.repeat(60, axis=1)
+        gripper_pts = gripper_pts.reshape(-1, 3)
+        gripper_trajectory = gripper_trajectory.reshape(-1, 4, 4)
+        gripper_pts_2_obj =  (gripper_trajectory[:, :3, :3] @ gripper_pts[..., None]).squeeze(-1)+ gripper_trajectory[:, :3, 3]
+        np.savetxt(save_path, gripper_pts_2_obj, fmt="%.6f")
+    
+    def axis_angle_to_matrix(self, axis, angle):
+
+        axis = axis / np.linalg.norm(axis)
+        K = np.array([[0, -axis[2], axis[1]],
+                      [axis[2], 0, -axis[0]],
+                      [-axis[1], axis[0], 0]])
+        R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)
+        T = np.eye(4)
+        T[:3, :3] = R
+        return T
+
+    def check_completion(self, objects):
+        if self.__len__() == 0:
+            return True
+
+        succ = False
+        if self.step_id >= 0:
+            succ = True
+
+        if succ:
+            self.step_id += 1
+
+        return succ

data_gen_dependencies/action/push.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, pose_difference
 
 
 class PushStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), vector_direction=None, extra_params=None, **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), vector_direction=None, extra_params=None, active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
 
         extra_params = {} if extra_params is None else extra_params
@@ -73,23 +73,5 @@ class PushStage(StageTemplate):
         return succ
 
 
-class ClickStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=None, vector_direction=None, **kwargs):
-        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
-        self.pre_distance = -0.03
-        self.move_distance = 0.00
-        self.generate_substage(target_pose, vector_direction)
 
-        
-    def generate_substage(self, target_pose, vector_direction):
-        vector_direction = vector_direction / np.linalg.norm(vector_direction)
-        # moveTo pre-place position
-        pre_pose = target_pose.copy()
-        pre_pose[:3,3] += vector_direction * self.pre_distance
-        self.sub_stages.append([pre_pose, 'close', np.eye(4), 'AvoidObs'])
-        
-        # insert
-        move_pose = target_pose.copy()
-        move_pose[:3,3] += vector_direction * self.move_distance
-        self.sub_stages.append([move_pose, None, np.eye(4), 'Simple'])
 

data_gen_dependencies/action/push_revolute.py

@@ -0,0 +1,146 @@
+import os
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+from data_gen_dependencies.action.base import StageTemplate
+
+
+class PushRevoluteStage(StageTemplate):
+    DELTA_DISTANCE = 0.003  # meter
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects=[], **kwargs):
+        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
+        #import ipdb;ipdb.set_trace()
+        self.objects = objects
+        self.passive_obj = objects[passive_obj_id]
+        self.active_obj = objects[active_obj_id]
+        main_part_name= passive_element["part_id"]
+        self.obj_base = os.path.dirname(passive_obj_id)
+        main_part_id = os.path.join(self.obj_base, main_part_name)
+        if main_part_id not in objects:
+            raise ValueError(f"Cannot find main part {main_part_name} in objects")
+        self.root_id = os.path.join(self.obj_base, "root")
+        self.main_part_obj = objects[main_part_id]
+        self.root_obj = objects[self.root_id]
+        self.revolute_radius = np.linalg.norm(self.passive_obj.obj_pose[:3,3] - self.main_part_obj.obj_pose[:3,3])
+        self.revolute_joint_position = self.main_part_obj.obj_pose[:3,3]
+        self.joint_position_threshold = passive_element.get('joint_position_threshold', 0.7)
+        self.correspond_joint_id = passive_element.get('correspond_joint_id', None)
+        correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id]
+        
+        self.joint_lower_limit = correspond_joint_info["lower_bound"]
+        self.joint_upper_limit = correspond_joint_info["upper_bound"]
+        self.joint_axis = correspond_joint_info["joint_axis"]
+        self.joint_type = correspond_joint_info["joint_type"]
+        assert self.joint_type == 'revolute', "joint_type must be revolute for push_revolute action"
+        if self.joint_position_threshold is None:
+            self.joint_position_threshold = 0.7
+        assert self.joint_position_threshold >= 0 and self.joint_position_threshold <= 1 
+
+        self.joint_direction = passive_element.get('joint_direction', 1)
+        assert self.joint_direction in [-1, 1]
+
+        self.joint_velocity_threshold = passive_element.get('joint_velocity_threshold', 999)
+        vector_direction = passive_element['direction']
+
+        self.extra_params = {} if extra_params is None else extra_params
+        self.generate_substage(target_pose, vector_direction)
+
+    def generate_substage(self, gripper2obj, vector_direction):
+        vector_direction = vector_direction / np.linalg.norm(vector_direction)
+        total_angle = self.joint_position_threshold * (self.joint_upper_limit - self.joint_lower_limit)
+        arc_length = abs(self.revolute_radius * total_angle)
+        step = int(arc_length / self.DELTA_DISTANCE)
+        from pyboot.utils.log import Log
+        Log.debug(f'PushRevoluteStage: total_angle={total_angle}, arc_length={arc_length}, step={step}')
+        trajectory = self.generate_trajectory(gripper2obj, vector_direction)
+        self.save_visualized_gripper_trajectory(trajectory, "push_revolute_traj.txt")
+
+        self.sub_stages.append([trajectory, None, np.eye(4), 'Trajectory'])
+        self.sub_stages.append([None, "open", np.eye(4), 'Simple'])
+
+        free_delta_pose = np.eye(4)
+        free_delta_pose[2,3] = -self.DELTA_DISTANCE
+        self.sub_stages.append([free_delta_pose, None, np.eye(4), 'local_gripper'])
+
+
+
+    def generate_trajectory(self, gripper2passive_obj, init_vector_direction):
+        trajectory = []
+        total_angle = self.joint_position_threshold * (self.joint_upper_limit - self.joint_lower_limit)
+        total_angle *= self.joint_direction
+
+        arc_length = abs(self.revolute_radius * total_angle)
+        num_steps = int(np.ceil(arc_length / self.DELTA_DISTANCE))
+        if num_steps < 1:
+            num_steps = 1
+
+        delta_theta = total_angle / num_steps
+        rotation_axis = np.array(self.joint_axis, dtype=float)
+        rotation_axis /= np.linalg.norm(rotation_axis) 
+
+        passive_obj_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.passive_obj.obj_pose
+        gripper_2_root = passive_obj_2_root @ gripper2passive_obj
+        pos0 = gripper_2_root[:3, 3]
+        R0 = gripper_2_root[:3, :3]
+
+        joint_2_root = np.linalg.inv(self.root_obj.obj_pose) @ self.main_part_obj.obj_pose
+
+        pivot = joint_2_root[:3, 3]
+
+        for i in range(1, num_steps + 1):
+            angle = i * delta_theta
+            R_delta = R.from_rotvec(rotation_axis * angle).as_matrix()
+            pos_rot = R_delta @ (pos0 - pivot) + pivot
+            R_new = R_delta @ R0
+
+            waypoint_2_root = np.eye(4)
+            waypoint_2_root[:3, :3] = R_new
+            waypoint_2_root[:3, 3] = pos_rot
+            waypoint_2_obj = np.linalg.inv(passive_obj_2_root) @ waypoint_2_root
+            trajectory.append(waypoint_2_obj)
+        return np.asarray(trajectory)
+
+    def save_visualized_gripper_trajectory(self, gripper_trajectory, save_path):
+        #import ipdb; ipdb.set_trace()
+        N = len(gripper_trajectory)
+        gripper_pts = np.zeros((1, 60, 3))
+        for i in range(10):
+            gripper_pts[0, i] = np.array([0, 0.01, 0.002*i])
+            gripper_pts[0, i+10] = np.array([0, 0.01, 0.002*i + 0.001])
+            gripper_pts[0, i+20] = np.array([0, -0.01, 0.002*i])
+            gripper_pts[0, i+30] = np.array([0, -0.01, 0.002*i + 0.001])
+            gripper_pts[0, i+40] = np.array([0, -0.01 + 0.002 * i, 0])
+        for i in range(5):
+            gripper_pts[0, i+50] = np.array([0, 0, -0.002*i])
+            gripper_pts[0, i+55] = np.array([0, 0, -0.002*i + 0.001])
+        gripper_pts -= np.array([0, 0, 0.002*10 + 0.001])
+        gripper_pts = gripper_pts.repeat(N, axis=0)
+        gripper_trajectory = gripper_trajectory.reshape(N, 1, 4, 4)
+        gripper_trajectory = gripper_trajectory.repeat(60, axis=1)
+        gripper_pts = gripper_pts.reshape(-1, 3)
+        gripper_trajectory = gripper_trajectory.reshape(-1, 4, 4)
+        gripper_pts_2_obj =  (gripper_trajectory[:, :3, :3] @ gripper_pts[..., None]).squeeze(-1)+ gripper_trajectory[:, :3, 3]
+        np.savetxt(save_path, gripper_pts_2_obj, fmt="%.6f")
+
+    def axis_angle_to_matrix(self, axis, angle):
+
+        axis = axis / np.linalg.norm(axis)
+        K = np.array([[0, -axis[2], axis[1]],
+                      [axis[2], 0, -axis[0]],
+                      [-axis[1], axis[0], 0]])
+        R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)
+        T = np.eye(4)
+        T[:3, :3] = R
+        return T
+
+    def check_completion(self, objects):
+        if self.__len__() == 0:
+            return True
+
+        succ = False
+        if self.step_id >= 0:
+            succ = True
+
+        if succ:
+            self.step_id += 1
+
+        return succ

data_gen_dependencies/action/slide.py

@@ -4,7 +4,7 @@ from data_gen_dependencies.action.base import StageTemplate, pose_difference
 
 
 class SlideStage(StageTemplate):
-    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), **kwargs):
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), active_obj=None, passive_obj=None, **kwargs):
         super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
 
         self.pre_distance = 0.0
@@ -39,7 +39,7 @@ class SlideStage(StageTemplate):
         current_obj_pose = objects[self.passive_obj_id].obj_pose
         pos_diff, _ = pose_difference(current_obj_pose, target_obj_pose)
         succ = pos_diff < 0.04
-        print(pos_diff, succ)
+        #print(pos_diff, succ)
         if succ:
             self.step_id += 1
         return succ

data_gen_dependencies/action/twist.py

@@ -0,0 +1,66 @@
+import numpy as np
+
+from data_gen_dependencies.action.base import StageTemplate
+
+
+class TwistStage(StageTemplate):
+    DELTA_DISTANCE = 0.01  # meter
+    def __init__(self, active_obj_id, passive_obj_id, active_element=None, passive_element=None, target_pose=np.eye(4), extra_params=None, objects = [], **kwargs):
+        super().__init__(active_obj_id, passive_obj_id, active_element, passive_element)
+        self.passive_obj = objects[passive_obj_id]
+        self.active_obj = objects[active_obj_id]
+        self.joint_position_threshold = passive_element.get('joint_position_threshold', 0.7)
+        if self.joint_position_threshold is None:
+            self.joint_position_threshold = 0.7
+        self.correspond_joint_id = passive_element.get('correspond_joint_id', None)
+        correspond_joint_info = self.passive_obj.joints_info[self.correspond_joint_id]
+        self.joint_lower_limit = correspond_joint_info["lower_bound"]
+        self.joint_upper_limit = correspond_joint_info["upper_bound"]
+        if self.joint_lower_limit is None or self.joint_upper_limit is None:
+            self.twist_degree_range = np.pi
+        else:
+            self.twist_degree_range = abs(self.joint_upper_limit - self.joint_lower_limit)
+        self.joint_axis = correspond_joint_info["joint_axis"]
+        self.joint_type = correspond_joint_info["joint_type"]
+        assert self.joint_type == 'continuous', "joint_type must be continuous for twist action"
+        if self.joint_position_threshold is None:
+            self.joint_position_threshold = 0.7
+        assert self.joint_position_threshold >= 0 and self.joint_position_threshold <= 1 
+
+        self.joint_direction = passive_element.get('joint_direction', 1)
+        assert self.joint_direction in [-1, 1]
+
+        self.joint_velocity_threshold = passive_element.get('joint_velocity_threshold', 999)
+        vector_direction = passive_element['direction']
+
+        self.extra_params = {} if extra_params is None else extra_params
+        self.generate_substage(target_pose, vector_direction)
+
+    def generate_substage(self, target_pose, vector_direction):
+        vector_direction = vector_direction / np.linalg.norm(vector_direction)
+        delta_twist_pose = self.axis_angle_to_matrix(np.asarray([0,0,1]), self.twist_degree_range * self.joint_position_threshold * self.joint_direction)
+        self.sub_stages.append([delta_twist_pose, None, np.eye(4), 'local_gripper'])
+
+    def axis_angle_to_matrix(self, axis, angle):
+
+        axis = axis / np.linalg.norm(axis)
+        K = np.array([[0, -axis[2], axis[1]],
+                      [axis[2], 0, -axis[0]],
+                      [-axis[1], axis[0], 0]])
+        R = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * (K @ K)
+        T = np.eye(4)
+        T[:3, :3] = R
+        return T
+
+    def check_completion(self, objects):
+        if self.__len__() == 0:
+            return True
+
+        succ = False
+        if self.step_id >= 0:
+            succ = True
+
+        if succ:
+            self.step_id += 1
+
+        return succ

data_gen_dependencies/base_agent.py

@@ -4,7 +4,7 @@ import time
 import numpy as np
 from scipy.spatial.transform import Rotation as R, Slerp
 
-from data_gen_dependencies.robot import IsaacSimRpcRobot
+from data_gen_dependencies.omni_robot import IsaacSimRpcRobot
 
 
 class BaseAgent(object):
@@ -14,19 +14,6 @@ class BaseAgent(object):
 
     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"]
@@ -54,9 +41,6 @@ class BaseAgent(object):
         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,
@@ -75,9 +59,6 @@ class BaseAgent(object):
     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"]:
@@ -122,7 +103,7 @@ class BaseAgent(object):
         return True
 
     def start_recording(self, task_name, camera_prim_list):
-        print(camera_prim_list)
+        #print(camera_prim_list)
         self.robot.client.start_recording(
             task_name=task_name,
             fps=30,

data_gen_dependencies/client.py

@@ -30,7 +30,7 @@ import time
 import json
 import pinocchio
 import os
-
+from pyboot.utils.log import Log
 
 # 目前代码中所有的旋转角度都以角度为单位
 class Rpc_Client:
@@ -43,12 +43,12 @@ class Rpc_Client:
                 grpc.channel_ready_future(self.channel).result(timeout=5)
                 self.robot_urdf = robot_urdf
             except grpc.FutureTimeoutError as e:
-                print(f"Failed to connect to gRPC server[{i}] ({client_host}): {e}", flush=True)
+                Log.warning(f"Failed to connect to gRPC server[{i}] ({client_host}): {e}")
                 time.sleep(3)
                 if i >= 599:
                     raise e
             except grpc.RpcError as e:
-                print(f"Failed to connect to gRPC server[{i}] ({client_host}): {e}", flush=True)
+                Log.warning(f"Failed to connect to gRPC server[{i}] ({client_host}): {e}")
                 time.sleep(3)
                 if i >= 599:
                     raise e
@@ -71,7 +71,7 @@ class Rpc_Client:
             "target_pose": target_pose
         }
         _frame_state: str = json.dumps(frame_state)
-        print(_frame_state)
+        #print(_frame_state)
         req.frame_state = _frame_state
         response = stub.SetFrameState(req)
         return response
@@ -423,7 +423,7 @@ class Rpc_Client:
             else:
                 joint_positions.append(joint_data[name])
         joint_positions = np.array(joint_positions)
-        print(joint_positions)
+        #print(joint_positions)
         pinocchio.forwardKinematics(model, data, joint_positions)
         J = pinocchio.computeFrameJacobian(model, data, joint_positions, joint_id)
         manip = np.sqrt(np.linalg.det(np.dot(J, J.T)))
@@ -552,7 +552,7 @@ class Rpc_Client:
         stub = sim_observation_service_pb2_grpc.SimObservationServiceStub(self.channel)
         req = sim_observation_service_pb2.SetMaterailReq()
         for mat in material_info:
-            print(mat)
+            #print(mat)
             mat_info = sim_observation_service_pb2.MaterialInfo()
             mat_info.object_prim = mat["object_prim"]
             mat_info.material_name = mat["material_name"]
@@ -567,7 +567,7 @@ class Rpc_Client:
         stub = sim_observation_service_pb2_grpc.SimObservationServiceStub(self.channel)
         req = sim_observation_service_pb2.SetLightReq()
         for light in light_info:
-            print(light)
+            #print(light)
             light_cfg = sim_observation_service_pb2.LightCfg()
             light_cfg.light_type = light["light_type"]
             light_cfg.light_prim = light["light_prim"]
@@ -579,384 +579,3 @@ class Rpc_Client:
             req.lights.append(light_cfg)
         response = stub.SetLight(req)
         return (response)
-
-
-# 调用示例
-def run():
-    rpc_client = Rpc_Client(client_host="localhost:50051")
-    try:
-        while True:
-            send_msg = input("input Command: ")
-            if send_msg == "30":
-                result = rpc_client.DrawLine(
-                    point_list_1=[[np.random.uniform(-5, 5), np.random.uniform(-5, 5), np.random.uniform(-5, 5)]],
-                    point_list_2=[[np.random.uniform(-5, 5), np.random.uniform(-5, 5), np.random.uniform(-5, 5)]],
-                    colors=[[np.random.uniform(0, 1), np.random.uniform(0, 1), np.random.uniform(0, 1)]],
-                    sizes=[np.random.uniform(0, 1)],
-                    name="test"
-                )
-            if send_msg == "3111":
-                result = rpc_client.ClearLine(name="test")
-            if send_msg == "point":
-                rpc_client.SetTargetPoint([0, 1, 1])
-            if send_msg == "test_mat":
-                rpc_client.SetMaterial([{
-                    "object_prim": "/World/huxing/a_keting/yangbanjian_A_016",
-                    "material_name": "Ash",
-                    "material_path": "materials/wood/Ash"
-                },
-                    {
-                        "object_prim": "/World/huxing/a_Base/Floor_003",
-                        "material_name": "Ash",
-                        "material_path": "materials/wood/Ash"
-
-                    }])
-                rpc_client.SetLight([{
-                    "light_type": "Distant",
-                    "light_prim": "/World/huxing/DistantLight",
-                    "light_temperature": 2000,
-                    "light_intensity": 1000,
-                    "rotation": [1, 0.5, 0.5, 0.5],
-                    "texture": ""
-                },
-                    {
-                        "light_type": "Dome",
-                        "light_prim": "/World/DomeLight",
-                        "light_temperature": 6500,
-                        "light_intensity": 1000,
-                        "rotation": [1, 0, 0, 0],
-                        "texture": "materials/hdri/abandoned_hall_01_4k.hdr"
-                    },
-                    {
-                        "light_type": "Rect",
-                        "light_prim": "/World/RectLight",
-                        "light_temperature": 4500,
-                        "light_intensity": 1000,
-                        "rotation": [1, 0, 0, 0],
-                        "texture": ""
-                    },
-                    {
-                        "light_type": "Rect",
-                        "light_prim": "/World/RectLight_01",
-                        "light_temperature": 8500,
-                        "light_intensity": 1000,
-                        "rotation": [1, 0, 0, 0],
-                        "texture": ""
-                    }])
-
-            if send_msg == "908":
-                pose_1 = [1, 1, 1], [1, 0, 0, 0]
-                pose_2 = [1, 0, 1], [1, 0, 0, 0]
-                trajecory_list = [pose_1, pose_2] * 10
-                result = rpc_client.SetTrajectoryList(trajecory_list, True)
-            if send_msg == "31":
-                target_poses = []
-                for i in range(1):
-                    # x = np.random.uniform(0.3, 0.7)
-                    # y = np.random.uniform(-0.3, 0.3)
-                    # z = np.random.uniform(0.96, 1.2)
-                    # pose = {
-                    #     "position": [x,y,z],
-                    #     "rotation": [1,0,0,0]
-                    # }
-                    x = 0.57597359649470348
-                    y = -0.45669529659303565
-                    z = 1.0198275517174573
-                    rw = 0.066194084385260935
-                    rx = 0.70713063274436749
-                    ry = 0.70071350186850612
-                    rz = 0.0677141028599091
-                    pose = {
-                        "position": [x, y, z],
-                        "rotation": [rw, rx, ry, rz]
-                    }
-                    target_poses.append(pose)
-
-                result = rpc_client.GetIKStatus(target_poses=target_poses, is_right=False)
-                print(result)
-            if send_msg == "1":
-                # 相机
-                # A2W
-                rpc_client.capture_frame(camera_prim_path="/World/Raise_A2/base_link/Head_Camera")
-                result = rpc_client.capture_semantic_frame(camera_prim_path="/World/Raise_A2/base_link/Head_Camera")
-                print(result)
-            if send_msg == "2":
-                # 臂
-                x = np.random.uniform(0.3, 0.7)
-                y = np.random.uniform(-0.3, 0)
-                z = np.random.uniform(0.1, 0.5)
-                rpc_client.moveto(target_position=[x, y, z], target_quaternion=np.array([0.0, -0.0, -1.0, 0.0]),
-                                  arm_name="left", is_backend=True, ee_interpolation=True)
-            if send_msg == "20":
-                result = rpc_client.MultiPlan(robot_name="left", target_poses=[
-                    [[0.5169683509992052, 0.1313259611510117, 1.1018942820728093],
-                     [0.40020943, 0.57116637, 0.69704651, -0.16651593]],
-                    [[0.5610938560120418, 0.048608636026916924, 1.0269891277236924],
-                     [0.40020943, 0.57116637, 0.69704651, -0.16651593]],
-                    [[0.5610938560120418, 0.048608636026916924, 1.2269891277236924],
-                     [0.40020943, 0.57116637, 0.69704651, -0.16651593]]
-                ])
-                print(result[0])
-
-            if send_msg == "44":
-                import time
-                # test cabinet
-                rpc_client.InitRobot("A2D_120s_horizon.json", "A2D.usd", "omnimanip_Simple_Room_01/simple_room.usd",
-                                     init_position=[-0.4, 0, -0.55])
-                rpc_client.reset()
-                time.sleep(1)
-                rpc_client.DetachObj()
-                rpc_client.reset()
-                time.sleep(1)
-                rpc_client.add_object(
-                    usd_path="objects/guanglun/storagefurniture/storagefurniture011/Storagefurniture011.usd",
-                    prim_path="/World/Objects/storagefurniture", label_name="test_storagefurniture",
-                    target_position=np.array([
-                        0.64,
-                        -0.3,
-                        0.40]), target_quaternion=np.array([0., 0., 0.70711, 0.70711]),
-                    target_scale=np.array([1, 1, 1]), color=np.array([1, 0, 1]),
-                    material="Plastic", mass=0.01)
-                time.sleep(1)
-                rpc_client.set_gripper_state(gripper_command="open", is_right=True, opened_width=1)
-
-                target_poses = []
-                start_time = time.time()
-                for i in range(50):
-                    x = np.random.uniform(0, 0.2)
-                    y = np.random.uniform(-0.3, 0.3)
-                    z = np.random.uniform(0.96, 1.2)
-                    pose = {
-                        "position": [x, y, z],
-                        "rotation": [1, x, y, z]
-                    }
-                    target_poses.append(pose)
-                result = rpc_client.GetIKStatus(target_poses=target_poses, is_right=False)
-                result = rpc_client.GetIKStatus(target_poses=target_poses, is_right=True)
-                rpc_client.DetachObj()
-                print("open")
-                rpc_client.get_object_pose("/World/Objects/storagefurniture/Storagefurniture011_Handle1")
-                rpc_client.get_object_pose("/World/Objects/storagefurniture/Storagefurniture011_Handle2")
-                obj_pose = rpc_client.get_object_pose("/World/Objects/storagefurniture/Storagefurniture011_Handle3")
-                obj_xyz = np.array(
-                    [obj_pose.object_pose.position.x, obj_pose.object_pose.position.y, obj_pose.object_pose.position.z])
-                ee_pose = rpc_client.GetEEPose(is_right=True)
-                ee_quat = np.array(
-                    [ee_pose.ee_pose.rpy.rw, ee_pose.ee_pose.rpy.rx, ee_pose.ee_pose.rpy.ry, ee_pose.ee_pose.rpy.rz])
-                goal_xyz = obj_xyz.copy()
-                goal_xyz[0] = goal_xyz[0] + 0.3
-                goal_xyz[2] = goal_xyz[2] + 0.55
-                print(goal_xyz)
-                rpc_client.moveto(
-                    target_position=goal_xyz,
-                    target_quaternion=ee_quat,
-                    arm_name="left",
-                    is_backend=False,
-                    ee_interpolation=False,
-                    distance_frame=0.001)
-                obj_pose = rpc_client.get_object_pose("/World/Objects/storagefurniture/Storagefurniture011_Handle3")
-                obj_xyz = np.array(
-                    [obj_pose.object_pose.position.x, obj_pose.object_pose.position.y, obj_pose.object_pose.position.z])
-                ee_pose = rpc_client.GetEEPose(is_right=True)
-                ee_quat = np.array(
-                    [ee_pose.ee_pose.rpy.rw, ee_pose.ee_pose.rpy.rx, ee_pose.ee_pose.rpy.ry, ee_pose.ee_pose.rpy.rz])
-                goal_xyz = obj_xyz.copy()
-                rpc_client.get_object_pose("/World/Objects/storagefurniture/Storagefurniture011_Handle1")
-                rpc_client.get_object_pose("/World/Objects/storagefurniture/Storagefurniture011_Handle2")
-                rpc_client.moveto(
-                    target_position=goal_xyz,
-                    target_quaternion=ee_quat,
-                    arm_name="left",
-                    is_backend=True,
-                    ee_interpolation=False,
-                    distance_frame=0.001)
-                rpc_client.set_gripper_state(gripper_command="close", is_right=True, opened_width=0.02)
-                time.sleep(1)
-                obj_pose = rpc_client.get_object_pose("/World/Objects/storagefurniture/Storagefurniture011_Handle3")
-                obj_xyz = np.array(
-                    [obj_pose.object_pose.position.x, obj_pose.object_pose.position.y, obj_pose.object_pose.position.z])
-                ee_pose = rpc_client.GetEEPose(is_right=True)
-                ee_quat = np.array(
-                    [ee_pose.ee_pose.rpy.rw, ee_pose.ee_pose.rpy.rx, ee_pose.ee_pose.rpy.ry, ee_pose.ee_pose.rpy.rz])
-                goal_xyz = obj_xyz.copy()
-                goal_xyz[0] = goal_xyz[0] - 0.1
-                rpc_client.moveto(
-                    target_position=goal_xyz,
-                    target_quaternion=ee_quat,
-                    arm_name="left",
-                    is_backend=True,
-                    ee_interpolation=False,
-                    distance_frame=0.001)
-            if send_msg == "21":
-                rpc_client.MultiMove(robot_name="left", plan_index=0)
-            if send_msg == "22":
-                rpc_client.MultiMove(robot_name="left", plan_index=1)
-            if send_msg == "23":
-                rpc_client.MultiMove(robot_name="left", plan_index=2)
-            if send_msg == "24":
-                rpc_client.start_recording()
-
-            if send_msg == "3":
-                # 关节
-                joint_states = rpc_client.get_joint_positions().states
-                joint_datas = {}
-                for joint in joint_states:
-                    joint_datas[joint.name] = joint.position
-                result = rpc_client.GetManipulability(joint_datas)
-                print(result)
-            if send_msg == "4":
-                # 手
-                rpc_client.set_gripper_state(gripper_command="open", is_right=True, opened_width=0.1)
-            if send_msg == "5":
-                # 物体
-
-                import time
-                # rpc_client.reset()
-                rpc_client.add_object(usd_path="Collected_cabinet_000/cabinet_000.usd",
-                                      prim_path="/World/Objects/cabinet", label_name="test_cabinet",
-                                      target_position=np.array([.8, 0., -0.1]),
-                                      target_quaternion=np.array([0., 0., 0.70711, 0.70711]),
-                                      target_scale=np.array([1, 1, 1]), color=np.array([1, 0, 1]),
-                                      material="Plastic", mass=0.01)
-                time.sleep(1)
-                result = rpc_client.get_object_joint("/World/Objects/cabinet")
-                print(result)
-                # x = np.random.uniform(0.2, 0.64)
-                # y = np.random.uniform(-0.3, 0.3)
-                # rpc_client.add_object(usd_path="genie3D/09.usd",
-                #                     prim_path="/World/Objects/object8", label_name="test7",
-                #                     target_position=np.array([np.random.uniform(0.2, 0.64),np.random.uniform(-0.3, 0.3),1.2]), target_quaternion=np.array([0.0,-0.0,-1.0, 0.0]), target_scale=np.array([0.01,0.01,0.01]), color=np.array([1,1,0]), material="Brass", mass=1)
-
-                # rpc_client.add_object(usd_path="genie3D/02/02.usd",
-                #                     prim_path="/World/Objects/object2", label_name="test2",
-                #                     target_position=np.array([np.random.uniform(0.2, 0.64), np.random.uniform(-0.3, 0.3),1.2]), target_quaternion=np.array([0.0,-0.0,-1.0, 0.0]), target_scale=np.array([0.01,0.01,0.01]), color=np.array([1,1,1]))
-                # rpc_client.get_object_pose(prim_path="/World/Xform/box_place")
-
-                # rpc_client.add_object(usd_path="genie3D/06/06.usd",
-                #                     prim_path="/World/Objects/object6", label_name="test6",
-                #                     target_position=np.array([np.random.uniform(0.2, 0.64), np.random.uniform(-0.3, 0.3),1.2]), target_quaternion=np.array([0.0,-0.0,-1.0, 0.0]), target_scale=np.array([0.01,0.01,0.01]), color=np.array([1,0,0]))
-                # rpc_client.get_object_pose(prim_path="/World/Xform/box_pick_00")
-
-                # rpc_client.add_object(usd_path="genie3D/01/01.usd",
-                #                     prim_path="/World/Objects/object1", label_name="test1",
-                #                     target_position=np.array([np.random.uniform(0.2, 0.64), np.random.uniform(-0.3, 0.3),1.2]), target_quaternion=np.array([0.0,-0.0,-1.0, 0.0]), target_scale=np.array([0.01,0.01,0.01]), color=np.array([0,0,1]))
-                # result = rpc_client.get_object_pose(prim_path="/World/Raise_A2_W_T1/head_link/visuals")
-            if send_msg == "6":
-                # 获取全部信息
-                rpc_client.get_joint_positions()
-            if send_msg == "7":
-                rpc_client.set_gripper_state(gripper_command="close", is_right=True, opened_width=0.02)
-            if send_msg == "8":
-                result = rpc_client.get_observation(
-                    data_keys={
-                        'camera': {
-                            'camera_prim_list': [
-                                '/World/Raise_A2/base_link/Head_Camera'
-                            ],
-                            'render_depth': True,
-                            'render_semantic': True
-                        },
-                        'pose': [
-                            '/World/Raise_A2/base_link/Head_Camera'
-                        ],
-                        'joint_position': True,
-                        'gripper': True
-                    }
-                )
-                print(result["camera"])
-            if send_msg == "9":
-                rpc_client.start_recording(data_keys={
-                    'camera': {
-                        'camera_prim_list': [
-                            '/World/Raise_A2/base_link/Head_Camera'
-                        ],
-                        'render_depth': True,
-                        'render_semantic': True
-                    },
-                    'pose': [
-                        '/World/Raise_A2/base_link/Head_Camera'
-                    ],
-                    'joint_position': True,
-                    'gripper': True
-                }, fps=30, task_name="test"
-
-                )
-                time.sleep(1)
-                rpc_client.stop_recording()
-                rpc_client.SendTaskStatus(False)
-            if send_msg == "10":
-                rpc_client.stop_recording()
-            if send_msg == "11":
-                rpc_client.reset()
-            if send_msg == "13":
-                result = rpc_client.Exit()
-                print(result)
-
-            if send_msg == '112':
-                rpc_client.InitRobot(robot_cfg="Franka.json", robot_usd="Franka/franka.usd",
-                                     scene_usd="Pick_Place_Franka_Yellow_Table.usd")
-
-            if send_msg == '113':
-                position = [1.2, 0., 1.2]
-                rotation = [0.61237, 0.35355, 0.35355, 0.61237]
-                rotation = [0.65328, 0.2706, 0.2706, 0.65328]
-                width = 640
-                height = 480
-                rpc_client.AddCamera("/World/Sensors/Head_Camera", position, rotation, width, height, 18.14756, 20.955,
-                                     15.2908, False)
-                response = rpc_client.capture_frame(camera_prim_path="/World/Sensors/Head_Camera")
-                # cam_info
-                cam_info = {
-                    "W": response.color_info.width,
-                    "H": response.color_info.height,
-                    "K": np.array(
-                        [
-                            [response.color_info.fx, 0, response.color_info.ppx],
-                            [0, response.color_info.fy, response.color_info.ppy],
-                            [0, 0, 1],
-                        ]
-                    ),
-                    "scale": 1,
-                }
-                # rgb
-                # rgb = np.frombuffer(response.color_image.data, dtype=np.uint8).reshape(cam_info["H"], cam_info["W"], 4)[
-                #     :, :, :3
-                # ]
-                # import cv2
-                # cv2.imwrite("head_img.png", rgb)
-
-            if send_msg == '114':
-                position = [0.05, 0., 0.]
-                rotation = [0.06163, 0.70442, 0.70442, 0.06163]
-                width = 640
-                height = 480
-                rpc_client.AddCamera("/panda/panda_hand/Hand_Camera_1", position, rotation, width, height, 18.14756,
-                                     20.955, 15.2908, True)
-                response = rpc_client.capture_frame(camera_prim_path="/panda/panda_hand/Hand_Camera_1")
-                # cam_info
-                cam_info = {
-                    "W": response.color_info.width,
-                    "H": response.color_info.height,
-                    "K": np.array(
-                        [
-                            [response.color_info.fx, 0, response.color_info.ppx],
-                            [0, response.color_info.fy, response.color_info.ppy],
-                            [0, 0, 1],
-                        ]
-                    ),
-                    "scale": 1,
-                }
-                # rgb
-                rgb = np.frombuffer(response.color_image.data, dtype=np.uint8).reshape(cam_info["H"], cam_info["W"], 4)[
-                    :, :, :3
-                ]
-                import cv2
-                cv2.imwrite("hand_img.png", rgb)
-
-    except Exception as e:
-        print("failed.{}".format(e))
-        return False
-
-
-if __name__ == '__main__':
-    run()

data_gen_dependencies/data_generate.py

@@ -1,15 +1,16 @@
 import json
-
+from pyboot.utils.log import Log
 from data_gen_dependencies.omni_robot import IsaacSimRpcRobot
 from data_gen_dependencies.omniagent import Agent
 
-def generate_data(task_file, client_host, use_recording=True):
+def generate_data(task_list, client_host, use_recording=True):
+    task_file = task_list[0]
     task = json.load(open(task_file))
     robot_cfg = task["robot"]["robot_cfg"]
     stand = task["robot"]["stand"]
     robot_position = task["robot"]["init_position"]
     robot_rotation = task["robot"]["init_rotation"]
-    scene_usd = task["scene_usd"]
+    scene_usd = task["scene"]["scene_usd"]
 
     robot = IsaacSimRpcRobot(robot_cfg=robot_cfg, scene_usd=scene_usd,
                              client_host=client_host,
@@ -19,12 +20,13 @@ def generate_data(task_file, client_host, use_recording=True):
     render_semantic = False
     if "render_semantic" in task["recording_setting"]:
         render_semantic = task["recording_setting"]["render_semantic"]
-    agent.run(task_file=task_file,
+    agent.run(task_list=task_list,
               camera_list=task["recording_setting"]["camera_list"],
               use_recording=use_recording,
               workspaces=task['scene']['function_space_objects'],
               fps=task["recording_setting"]["fps"],
               render_semantic=render_semantic,
               )
-    print("job done")
+
+    Log.success("finish generating data")
     robot.client.Exit()

data_gen_dependencies/manip_solver.py

@@ -2,8 +2,11 @@
 import copy
 import json
 import os
+from sqlite3 import dbapi2
+from this import d
 import time
 
+from pyboot.utils.log import Log
 import trimesh
 
 curPath = os.path.abspath(os.path.dirname(__file__))
@@ -210,6 +213,7 @@ def get_aligned_pose(active_obj, passive_obj, distance=0.01, N=1):
 
 
 def load_task_solution(task_info):
+    Log.info('Loading task solution')
     # task_info = json.load(open('%s/task_info.json'%task_dir, 'rb'))
     stages = task_info['stages']
 
@@ -217,7 +221,8 @@ def load_task_solution(task_info):
         'gripper': OmniObject('gripper')
     }
 
-    for obj_info in task_info['objects']:
+    for i, obj_info in enumerate(task_info['objects']):
+        Log.info(f'Loading object: {obj_info["object_id"]} [Loading task solution: {i+1}/{len(task_info["objects"])}]')
         obj_id = obj_info['object_id']
         if obj_id == 'fix_pose':
             obj = OmniObject('fix_pose')
@@ -243,21 +248,30 @@ def load_task_solution(task_info):
             obj_dir = obj_info["data_info_dir"]
             obj = OmniObject.from_obj_dir(obj_dir, obj_info=obj_info)
             objects[obj_id] = obj
-
-            if hasattr(obj, 'part_ids'):
-                if hasattr(obj, 'part_joint_limits') and obj.part_joint_limits is not None:
-                    obj_parts_joint_limits = obj.part_joint_limits
-                for part_id in obj.part_ids:
+            if obj.is_articulated:
+                for part_id in obj.parts_info:
                     id = obj_id + '/%s' % part_id
                     objects[id] = copy.deepcopy(obj)
                     objects[id].name = id
-                    objects[id].part_joint_limit = obj_parts_joint_limits[part_id]
-                if len(obj.part_ids):
+                    objects[id].size = obj.parts_info[part_id]['size']
+                    joint_name = obj.parts_info[part_id]['correspond_joint_id']
+                    objects[id].joint_name = joint_name
+                    if objects[id].joint_name is not None and joint_name in obj.joints_info:
+                        objects[id].joint_type = obj.joints_info[joint_name]['joint_type']
+                        objects[id].joint_axis = obj.joints_info[joint_name]['joint_axis']
+                        objects[id].lower_bound = obj.joints_info[joint_name]['lower_bound']
+                        objects[id].upper_bound = obj.joints_info[joint_name]['upper_bound']
+      
+                if len(obj.parts_info):
                     del objects[obj_id]
+     
+    Log.success('Finished Loading task solution')
+
     return stages, objects
 
 
 def parse_stage(stage, objects):
+
     action = stage['action']
     if action in ['reset']:
         return action, 'gripper', 'gripper', None, None, None, None
@@ -272,24 +286,38 @@ def parse_stage(stage, objects):
     active_obj = objects[active_obj_id]
     passive_obj = objects[passive_obj_id]
 
-    single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe']
+    single_obj = action in ['pull', 'rotate', 'slide', 'shave', 'brush', 'wipe', 'pull_revolute', "twist"]
 
     def _load_element(obj, type):
+
         if action in ['pick', 'hook']:
             action_mapped = 'grasp'
+        elif action in ['pull_revolute']:
+            action_mapped = 'pull'
+        elif action in ['push_revolute']:
+            action_mapped = 'push'
+        elif action in ["press_prismatic"]:
+            action_mapped = 'press'
         else:
             action_mapped = action
-        if action_mapped == 'grasp' and type == 'active':
+        if (action_mapped == 'grasp' or action_mapped == 'pull') and type == 'active':
             return None, None
         elif obj.name == 'gripper':
             element = obj.elements[type][action_mapped]
             return element, 'default'
         primitive = stage[type]['primitive'] if stage[type]['primitive'] is not None else 'default'
         if primitive != 'default' or (action_mapped == 'grasp' and type == 'passive'):
-            if action_mapped not in obj.elements[type]:
+            if action_mapped not in obj.elements[type] and action not in obj.elements[type]:
+                #======= DEBUG START =======
+                print(f'debug: ')
+                import ipdb;ipdb.set_trace()
+                #=======  DEBUG END  =======
                 print('No %s element for %s' % (action_mapped, obj.name))
                 return None, None
-            element = obj.elements[type][action_mapped][primitive]
+            elif action_mapped in obj.elements[type]:
+                element = obj.elements[type][action_mapped][primitive]
+            elif action in obj.elements[type]:
+                element = obj.elements[type][action][primitive]
         else:
             element = []
             for primitive in obj.elements[type][action_mapped]:
@@ -316,11 +344,10 @@ def select_obj(objects, stages, robot):
 
     ''' 初筛抓取pose,得到 grasp_poses_canonical, grasp_poses '''
     grasp_stage_id = None
-
-    if stages[0]['action'] in ['push', 'reset', 'click']:
+    if stages[0]['action'] in ['push', 'reset', 'press_prismatic']:
         gripper2obj = current_gripper_pose
 
-    elif stages[0]['action'] in ['pick', 'grasp', 'hook']:
+    elif stages[0]['action'] in ['pick', 'grasp', 'hook', "pull_revolute", "twist",  "push_revolute"]:
         action = stages[0]['action']
 
         ''' 筛掉无IK解的grasp pose '''
@@ -328,6 +355,10 @@ def select_obj(objects, stages, robot):
         grasp_stage = parse_stage(stages[0], objects)
         _, _, passive_obj_id, _, passive_element, _, _ = grasp_stage
         grasp_obj_id = passive_obj_id
+        if stages[0]['action'] == "pull_revolute":
+            import ipdb;ipdb.set_trace()
+        if stages[0]['action'] == "push_revolute":
+            import ipdb;ipdb.set_trace()
         grasp_poses_canonical = passive_element['grasp_pose'].copy()
         grasp_widths = passive_element['width']
 
@@ -346,13 +377,10 @@ def select_obj(objects, stages, robot):
         ik_success, _ = robot.solve_ik(grasp_poses, ee_type='gripper', arm=arm, type='Simple')
         grasp_poses_canonical, grasp_poses = grasp_poses_canonical[ik_success], grasp_poses[ik_success]
         grasp_widths = grasp_widths[ik_success]
-        print('%s, %s, Filtered grasp pose with isaac-sim IK: %d/%d' % (action, passive_obj_id, grasp_poses.shape[0],
-                                                                        ik_success.shape[0]))
 
         if len(grasp_poses) == 0:
             print(action, 'No grasp_gripper_pose can pass Isaac IK')
             return []
-
     ''' 基于有IK解的grasp pose分数,选择最优的passive primitive element,同时选出最优的一个grasp pose'''
     if grasp_stage_id is not None:
         next_stage_id = grasp_stage_id + 1
@@ -380,7 +408,8 @@ def select_obj(objects, stages, robot):
                 direction_canonical = (np.linalg.inv(obj_pose) @ np.array([*direction, 0]))[:3]
                 active_elements = [{'xyz': xyz_canonical, 'direction': direction_canonical}]
 
-            # import ipdb; ipdb.set_trace()
+
+            #import ipdb; ipdb.set_trace()
             t0 = time.time()
             element_ik_score = []
             grasp_pose_ik_score = []
@@ -498,7 +527,6 @@ def split_grasp_stages(stages):
 
 def generate_action_stages(objects, all_stages, robot):
     split_stages = split_grasp_stages(all_stages)
-
     action_stages = []
     for stages in split_stages:
         gripper2obj = select_obj(objects, stages, robot)
@@ -506,30 +534,36 @@ def generate_action_stages(objects, all_stages, robot):
             print('No gripper2obj pose can pass IK')
             gripper2obj = select_obj(objects, stages, robot)
             return []
+        stage_idx = 0
         for stage in stages:
+            stage_idx += 1
             extra_params = stage.get('extra_params', {})
             arm = extra_params.get('arm', 'right')
-            current_gripper_pose = robot.get_ee_pose(id=arm)
             action, active_obj_id, passive_obj_id, active_elements, passive_elements, active_primitive, passive_primitive = parse_stage(
                 stage, objects)
             active_obj = objects[active_obj_id]
             passive_obj = objects[passive_obj_id]
-
-            single_obj = active_obj_id == passive_obj_id
+            Log.info(f'Generating action stage [{stage_idx}/{len(stages)}]: Action({action}), Active Object({active_obj_id}), Passive Object({passive_obj_id})')
 
             substages = None
+            #======= DEBUG START =======
+            print(f'debug: before build_stage {action}')
+            #import ipdb;ipdb.set_trace()
+            #=======  DEBUG END  =======
             if action in ['reset']:
                 substages = True
             elif action in ['pick', 'grasp', 'hook']:
                 substages = build_stage(action)(active_obj_id, passive_obj_id, active_elements, passive_elements,
-                                                gripper2obj, extra_params=stage.get('extra_params', None))
-            elif action in ['slide', 'shave', 'brush', 'wipe', 'pull', 'rotate', 'pull', 'move',
-                            'reset']:  # grasp + 物体自身运动
+                                                gripper2obj, extra_params=stage.get('extra_params', None), objects=objects)
+            elif action in ['slide', 'shave', 'brush', 'wipe', 'pull', 'rotate', 'move',
+                            'reset', "pull_revolute","push_revolute", "twist"]:  # grasp + 物体自身运动
                 passive_element = passive_elements[np.random.choice(len(passive_elements))]
                 substages = build_stage(action)(
                     active_obj_id=active_obj_id,
                     passive_obj_id=passive_obj_id,
-                    passive_element=passive_element
+                    passive_element=passive_element,
+                    objects=objects,
+                    target_pose=gripper2obj
                 )
             else:
                 passive_element = passive_elements[np.random.choice(len(passive_elements))]
@@ -569,7 +603,8 @@ def generate_action_stages(objects, all_stages, robot):
                     ik_jacobian_score = ik_info["jacobian_score"][ik_success]
 
                     if len(target_gripper_poses_pass_ik) == 0:
-                        print(action, ': No target_obj_pose can pass isaac curobo IK')
+                        Log.warning(f'{action}: No target_obj_pose can pass isaac curobo IK')
+                        import ipdb;ipdb.set_trace()
                         continue
 
                     target_joint_positions = []
@@ -611,12 +646,13 @@ def generate_action_stages(objects, all_stages, robot):
                         obj2part=obj2part,
                         vector_direction=passive_element['direction'],
                         passive_element=passive_element,
-                        extra_params=stage.get('extra_params', {})
+                        extra_params=stage.get('extra_params', {}),
+                        objects=objects
                     )
                     break
 
             if substages is None:
-                print(action, ': No target_obj_pose can pass IK')
+                Log.warning(f'{action}:substage is None')
                 return []
             action_stages.append((action, substages))
 

data_gen_dependencies/object.py

@@ -4,6 +4,7 @@ import pickle
 
 import numpy as np
 from grasp_nms import nms_grasp
+from pyboot.utils.log import Log
 from data_gen_dependencies.transforms import transform_coordinates_3d
 
 
@@ -48,9 +49,9 @@ class OmniObject:
                         "direction": np.array([-1, 0, 0.3]),
                     },
                 ],
-                'click': {
+                'press': {
                     "part": "finger edge",
-                    "task": "click",
+                    "task": "press",
                     "xyz": np.array([0, 0, 0]),
                     "direction": np.array([0, 0, 1]),
                 },
@@ -66,9 +67,9 @@ class OmniObject:
                     "xyz": np.array([0, 0, 0]),
                     "direction": np.array([0, 0, 0.08]),
                 },
-                'rotate': {
+                'twist': {
                     "part": "finger edge",
-                    "task": "rotate",
+                    "task": "twist",
                     "xyz": np.array([0, 0, 0]),
                     "direction": np.array([0, 0, 0.08]),
                 }
@@ -79,7 +80,6 @@ class OmniObject:
         if "interaction" in obj_info:
             obj_info = obj_info
             interaction_info = obj_info["interaction"]
-            part_joint_limits_info = obj_info.get("part_joint_limits", None)
 
         else:
 
@@ -92,7 +92,6 @@ class OmniObject:
             obj_info = json.load(open(obj_info_file))
             interaction_data = json.load(open(interaction_label_file))
             interaction_info = interaction_data['interaction']
-            part_joint_limits_info = interaction_data.get("part_joint_limits", None)
 
         obj = cls(
             name=obj_info['object_id'],
@@ -103,10 +102,9 @@ class OmniObject:
         if os.path.exists(mesh_file):
             obj_info['mesh_file'] = mesh_file
 
-        obj.part_joint_limits = part_joint_limits_info
 
         ''' Load interaction labels '''
-        obj.part_ids = []
+        
         for type in ['active', 'passive']:
             if type not in interaction_info:
                 continue
@@ -122,22 +120,24 @@ class OmniObject:
                         }
                         if isinstance(grasp_files, str):
                             grasp_files = [grasp_files]
+                        if isinstance(grasp_files, dict):
+                            grasp_files = [grasp_files['grasp_pose_file']]
                         for grasp_file in grasp_files:
                             grasp_file = '%s/%s' % (obj_dir, grasp_file)
                             if not os.path.exists(grasp_file):
-                                print('Grasp file %s not found' % grasp_file)
+                                Log.warning('Grasp file %s not found' % grasp_file)
                                 continue
                             _data = pickle.load(open(grasp_file, 'rb'))
                             _data['grasp_pose'] = np.array(_data['grasp_pose'])
                             _data['width'] = np.array(_data['width'])
 
                             if _data['grasp_pose'].shape[0] == 0:
-                                print('Empty grasp pose in %s' % grasp_file)
+                                Log.warning('Empty grasp pose in %s' % grasp_file)
                                 continue
                             grasp_data['grasp_pose'].append(_data['grasp_pose'])
                             grasp_data['width'].append(_data['width'])
                         if len(grasp_data['grasp_pose']) == 0:
-                            print('Empty grasp pose in %s' % grasp_file)
+                            Log.warning('Empty grasp pose in %s' % grasp_file)
                             continue
 
                         grasp_data['grasp_pose'] = np.concatenate(grasp_data['grasp_pose'])
@@ -178,30 +178,30 @@ class OmniObject:
                             grasp_data['grasp_pose'][:, :3, 3] = translation
                             grasp_data['width'] = width
 
-                            print(f"Grasp num of {obj.name} after NMS: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
+                            #print(f"Grasp num of {obj.name} after NMS: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
                             # downsample
                             if grasp_data['grasp_pose'].shape[0] > max_grasps:
                                 grasp_num = grasp_data['grasp_pose'].shape[0]
                                 index = np.random.choice(grasp_num, max_grasps, replace=False)
                                 grasp_data['grasp_pose'] = grasp_data['grasp_pose'][index]
                                 grasp_data['width'] = grasp_data['width'][index]
-                                print(
-                                    f"Grasp num of {obj.name} after downsample: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
+                                #print(
+                                 #   f"Grasp num of {obj.name} after downsample: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
                         else:
-                            print(f"Grasp num of {obj.name}: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
+                            pass
+                            #print(f"Grasp num of {obj.name}: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
 
                         action_info[grasp_part] = grasp_data
 
                     interaction_info[type][action] = action_info
-                else:
-                    for primitive in action_info:
-                        for primitive_info in action_info[primitive]:
-                            if 'part_id' in primitive_info:
-                                obj.part_ids.append(primitive_info['part_id'])
+        obj.is_articulated = obj_info.get("is_articulated", False)
+        if obj.is_articulated:
+            obj.parts_info = obj_info["parts_info"]
+            obj.joints_info = obj_info["joints_info"]
         obj.elements = interaction_info
         obj.info = obj_info
         # import ipdb;ipdb.set_trace()
-        obj.is_articulated = True if part_joint_limits_info is not None else False
+        
         return obj
 
     def set_element(self, element):

data_gen_dependencies/omni_robot.py

@@ -249,7 +249,7 @@ class IsaacSimRpcRobot(Robot):
         """
         type = content["type"]
         arm_name = content.get('arm', "right")
-        print(arm_name)
+        #print(arm_name)
         if type == "matrix":
             if isinstance(content["matrix"], list):
                 content["matrix"] = np.array(content["matrix"])
@@ -275,8 +275,8 @@ class IsaacSimRpcRobot(Robot):
                 target_rotation_matrix, target_position = target_matrix[:3, :3], target_matrix[:3, 3]
                 target_rotation = get_quaternion_from_euler(matrix_to_euler_angles(target_rotation_matrix), order="ZYX")
 
-                print(f"target_position is{target_position}")
-                print(f"target_rotation is{target_rotation}")
+                #print(f"target_position is{target_position}")
+                #print(f"target_rotation is{target_rotation}")
                 # import ipdb;ipdb.set_trace()
                 # inverse_local_orientation = [self.init_rotation[0], -1*self.init_rotation[1], -1*self.init_rotation[2], -1*self.init_rotation[3]]
                 # local_position = quat_multiplication(inverse_local_orientation,T-self.init_position)
@@ -291,7 +291,7 @@ class IsaacSimRpcRobot(Robot):
                     ).errmsg
                     == "True"
             )
-            print("move!", T, quat_wxyz, state)
+            #print("move!", T, quat_wxyz, state)
 
         elif type == "quat":
 

data_gen_dependencies/omniagent.py

@@ -1,17 +1,19 @@
 # -*- coding: utf-8 -*-
+
 import json
 import os
 import pickle
-import random
 import time
+import random
 
 import numpy as np
 from pyboot.utils.log import Log
 
-from data_gen_dependencies.transforms import calculate_rotation_matrix
-from data_gen_dependencies.manip_solver import load_task_solution, generate_action_stages, split_grasp_stages
 from data_gen_dependencies.base_agent import BaseAgent
+from data_gen_dependencies.manip_solver import load_task_solution, generate_action_stages, split_grasp_stages
 from data_gen_dependencies.omni_robot import IsaacSimRpcRobot
+from data_gen_dependencies.transforms import calculate_rotation_matrix
+
 
 class Agent(BaseAgent):
     def __init__(self, robot: IsaacSimRpcRobot):
@@ -35,6 +37,7 @@ class Agent(BaseAgent):
         )
 
     def generate_layout(self, task_file):
+        Log.info(f'Generating layout for task: {task_file}')
         self.task_file = task_file
         with open(task_file, "r") as f:
             task_info = json.load(f)
@@ -48,6 +51,7 @@ class Agent(BaseAgent):
                         task_info['objects'][i]['mass'] = 10
                         break
 
+        Log.info('Adding objects [Generating layout: 1/5]')
         self.articulated_objs = []
         for object_info in task_info["objects"]:
             if object_info['object_id'] == 'fix_pose':
@@ -59,8 +63,8 @@ class Agent(BaseAgent):
             self.add_object(object_info)
         time.sleep(2)
 
-        self.arm = task_info["arm"]
-
+        Log.info('Setting arm [Generating layout: 2/5]')
+        self.arm = task_info["robot"]["arm"]
         ''' For A2D: Fix camera rotaton to look at target object '''
         task_related_objs = []
         for stage in task_info['stages']:
@@ -78,24 +82,29 @@ class Agent(BaseAgent):
         target_lookat_point = np.mean(np.stack(target_lookat_point), axis=0)
         self.robot.client.SetTargetPoint(target_lookat_point.tolist())
 
+        #Log.info('Setting material [Generating layout: 3/5]')
+        Log.warning('Pass setting material [Generating layout: 3/5]')
         ''' Set material '''
-        material_infos = []
-        if "object_with_material" in task_info:
-            for key in task_info['object_with_material']:
-                material_infos += task_info['object_with_material'][key]
-            if len(material_infos):
-                self.robot.client.SetMaterial(material_infos)
-                time.sleep(0.3)
+        # material_infos = []
+        # if "object_with_material" in task_info:
+        #     for key in task_info['object_with_material']:
+        #         material_infos += task_info['object_with_material'][key]
+        #     if len(material_infos):
+        #         self.robot.client.SetMaterial(material_infos)
+        #         time.sleep(0.3)
 
+        #Log.info('Setting light [Generating layout: 4/5]')
+        Log.warning('Pass setting light [Generating layout: 4/5]')
         ''' Set light '''
-        light_infos = []
-        if "lights" in task_info:
-            for key in task_info['lights']:
-                light_infos += task_info['lights'][key]
-            if len(light_infos):
-                self.robot.client.SetLight(light_infos)
-                time.sleep(0.3)
+        # light_infos = []
+        # if "lights" in task_info:
+        #     for key in task_info['lights']:
+        #         light_infos += task_info['lights'][key]
+        #     if len(light_infos):
+        #         self.robot.client.SetLight(light_infos)
+        #         time.sleep(0.3)
 
+        Log.info('Setting camera [Generating layout: 5/5]')
         ''' Set camera'''
         if "cameras" in task_info:
             for cam_id in task_info['cameras']:
@@ -106,6 +115,7 @@ class Agent(BaseAgent):
                     cam_info['focal_length'], cam_info['horizontal_aperture'], cam_info['vertical_aperture'],
                     cam_info['is_local']
                 )
+        Log.success('Finished Generating Layout')
 
     def update_objects(self, objects, arm='right'):
         # update gripper pose
@@ -124,17 +134,16 @@ class Agent(BaseAgent):
                     objects['fix_pose'].obj_pose[:3, :3] = rotation_matrix
                 continue
                 # TODO(unify part_name and obj_name)
-            if '/' in obj_id:
+            obj = objects[obj_id]
+            if obj.is_articulated:
                 obj_name = obj_id.split('/')[0]
-                part_name = obj_id.split('/')[1]
-
                 object_joint_state = self.robot.client.get_object_joint('/World/Objects/%s' % obj_name)
                 for joint_name, joint_position, joint_velocity in zip(object_joint_state.joint_names,
                                                                       object_joint_state.joint_positions,
                                                                       object_joint_state.joint_velocities):
-                    if joint_name[-1] == part_name[-1]:
-                        objects[obj_id].joint_position = joint_position
-                        objects[obj_id].joint_velocity = joint_velocity
+                    if joint_name == obj.joint_name:
+                        obj.joint_position = joint_position
+                        obj.joint_velocity = joint_velocity
 
             objects[obj_id].obj_pose = self.robot.get_prim_world_pose('/World/Objects/%s' % obj_id)
             if hasattr(objects[obj_id], 'info') and 'simple_place' in objects[obj_id].info and objects[obj_id].info[
@@ -184,6 +193,8 @@ class Agent(BaseAgent):
                     file = objs_interaction[passive_obj_id]['passive']['grasp'][primitive]
                     if isinstance(file, list):
                         file = file[0]
+                    if isinstance(file, dict):
+                        file = file['grasp_pose_file']
                 grasp_file = os.path.join(data_root, obj_dir, file)
                 if not os.path.exists(grasp_file):
                     Log.error('-- Grasp file not exist: %s' % grasp_file)
@@ -194,143 +205,216 @@ class Agent(BaseAgent):
                     Log.error('-- Grasp file empty: %s' % grasp_file)
                     return False
         return True
+    
+    def visualize_gripper_pose(self, gripper_pose):
+        canonical_gripper_pts = np.zeros((60, 3))
+        for i in range(10):
+            canonical_gripper_pts[i] = np.array([0, 0.01, 0.002*i])
+            canonical_gripper_pts[i+10] = np.array([0, 0.01, 0.002*i + 0.001])
+            canonical_gripper_pts[i+20] = np.array([0, -0.01, 0.002*i])
+            canonical_gripper_pts[i+30] = np.array([0, -0.01, 0.002*i + 0.001])
+            canonical_gripper_pts[i+40] = np.array([0, -0.01 + 0.002 * i, 0])
+        for i in range(5):
+            canonical_gripper_pts[i+50] = np.array([0, 0, -0.002*i])
+            canonical_gripper_pts[i+55] = np.array([0, 0, -0.002*i + 0.001])
+        canonical_gripper_pts -= np.array([0, 0, 0.002*10 + 0.001])
+        canonical_gripper_pts = canonical_gripper_pts.reshape(-1, 3)
+        #simport ipdb; ipdb.set_trace()
+        gripper_pts =  (gripper_pose[:3, :3] @ canonical_gripper_pts.T).T + gripper_pose[:3, 3]
+        return gripper_pts
 
-    def run(self, task_file, camera_list, use_recording, workspaces, fps=10, render_semantic=False):
-        if not self.check_task_file(task_file):
-            Log.error("Task file bad: %s" % task_file)
-            return False
-        print("Start Task:", task_file)
-        self.reset()
-        self.attached_obj_id = None
+    def run(self, task_list, camera_list, use_recording, workspaces, fps=10, render_semantic=False):
+        for index, task_file in enumerate(task_list):
+            if not self.check_task_file(task_file):
+                Log.error("Task file bad: %s" % task_file)
+                continue
+            Log.info("Start task: "+ task_file)
+            self.reset()
+            self.attached_obj_id = None
 
-        # import ipdb;ipdb.set_trace()
+            # import ipdb;ipdb.set_trace()
 
-        self.generate_layout(task_file)
-        # import ipdb;ipdb.set_trace()
-        self.robot.open_gripper(id='right')
-        self.robot.open_gripper(id='left')
+            self.generate_layout(task_file)
+            # import ipdb;ipdb.set_trace()
+            self.robot.open_gripper(id='right')
+            self.robot.open_gripper(id='left')
 
-        self.robot.reset_pose = {
-            'right': self.robot.get_ee_pose(ee_type='gripper', id='right'),
-            'left': self.robot.get_ee_pose(ee_type='gripper', id='left'),
-        }
-        print('Reset pose:', self.robot.reset_pose)
+            self.robot.reset_pose = {
+                'right': self.robot.get_ee_pose(ee_type='gripper', id='right'),
+                'left': self.robot.get_ee_pose(ee_type='gripper', id='left'),
+            }
+            #print('Reset pose:', self.robot.reset_pose)
 
-        task_info = json.load(open(task_file, 'rb'))
-        stages, objects = load_task_solution(task_info)
-        objects = self.update_objects(objects)
-        split_stages = split_grasp_stages(stages)
-        # import ipdb; ipdb.set_trace()
-        if use_recording:
-            self.start_recording(task_name="[%s]" % (os.path.basename(task_file).split(".")[0]),
-                                 camera_prim_list=camera_list, fps=fps,
-                                 render_semantic=render_semantic)  # TODO 录制判断
+            task_info = json.load(open(task_file, 'rb'))
+            stages, objects = load_task_solution(task_info)
+            objects = self.update_objects(objects)
+            split_stages = split_grasp_stages(stages)
+            # import ipdb; ipdb.set_trace()
+            if use_recording:
+                self.start_recording(task_name="[%s]" % (os.path.basename(task_file).split(".")[0]),
+                                     camera_prim_list=camera_list, fps=fps,
+                                     render_semantic=render_semantic)  # TODO 录制判断
 
-        stage_id = -1
-        substages = None
-        for _stages in split_stages:
-            extra_params = _stages[0].get('extra_params', {})
-            active_id, passive_id = _stages[0]['active']['object_id'], _stages[0]['passive']['object_id']
-            arm = extra_params.get('arm', 'right')
-            action_stages = generate_action_stages(objects, _stages, self.robot)
-            if not len(action_stages):
-                success = False
-                print('No action stage generated.')
-                break
+            
+            success = False
+            substages = None
+            for _stages in split_stages:
+                stage_id = -1
+                extra_params = _stages[0].get('extra_params', {})
+                active_id, passive_id = _stages[0]['active']['object_id'], _stages[0]['passive']['object_id']
+                arm = extra_params.get('arm', 'right')
+                Log.info(f'Generating action stages')
+                action_stages = generate_action_stages(objects, _stages, self.robot)
+                if not len(action_stages):
+                    success = False
+                    print('No action stage generated.')
+                    break
 
-            # Execution
-            success = True
+                # Execution
+                success = True
+                last_gripper_state = None
+                for action, substages in action_stages:
+                    
+                    stage_id += 1
+                    Log.info(f'Start action stage: {action} [Executing task: {stage_id+1}/{len(action_stages)}]')
 
-            for action, substages in action_stages:
-                stage_id += 1
-                print('>>>>  Stage [%d]  <<<<' % (stage_id + 1))
-                active_id, passive_id = _stages[stage_id]['active']['object_id'], _stages[stage_id]['passive'][
-                    'object_id']
-                if action in ['reset']:
-                    init_pose = self.robot.reset_pose[arm]
-                    curr_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
-                    print(arm)
-                    print(curr_pose)
-                    interp_pose = init_pose.copy()
-                    interp_pose[:3, 3] = curr_pose[:3, 3] + (init_pose[:3, 3] - curr_pose[:3, 3]) * 0.25
-                    success = self.robot.move_pose(self.robot.reset_pose[arm], type='AvoidObs', arm=arm, block=True)
-                    continue
+                    active_id, passive_id = _stages[stage_id]['active']['object_id'], _stages[stage_id]['passive'][
+                        'object_id']
+                    
+                    if action in ['reset']:
+                        init_pose = self.robot.reset_pose[arm]
+                        curr_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
+                        interp_pose = init_pose.copy()
+                        interp_pose[:3, 3] = curr_pose[:3, 3] + (init_pose[:3, 3] - curr_pose[:3, 3]) * 0.25
+                        success = self.robot.move_pose(self.robot.reset_pose[arm], type='AvoidObs', arm=arm, block=True)
+                        continue
+                    substage_id = 0
+                    substage_len = len(substages)
+                    while len(substages):
+                        objects = self.update_objects(objects, arm=arm)
+                        substage_id += 1
+                        
+                        target_gripper_pose, motion_type, gripper_action, arm = substages.get_action(objects)
+                        target_gripper_pos = None if target_gripper_pose is None else target_gripper_pose[:3, 3]
+                        Log.info(f'-- ({action}) Sub-stage {substage_id}/{substage_len}: target_pos={target_gripper_pos}, motion_type={motion_type}, gripper_action={gripper_action}, arm={arm}')
+                        arm = extra_params.get('arm', 'right')
+                        self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
+                                                          self.attached_obj_id is not None, arm=arm,
+                                                          target_pose=target_gripper_pose)
+                        if target_gripper_pose is not None:
+                            self.robot.move_pose(target_gripper_pose, motion_type, arm=arm, block=True)
+                            # vis_gripper_pose = self.visualize_gripper_pose(target_gripper_pose)
+                            # import ipdb;ipdb.set_trace()
+                            # np.savetxt('debug_gripper.txt', vis_gripper_pose)
+                            
 
-                    # if action in ['grasp', 'pick']:
-                #     obj_id = substages.passive_obj_id
-                #     if  obj_id.split('/')[0] not in self.articulated_objs:
-                #         self.robot.target_object = substages.passive_obj_id
 
-                while len(substages):
-                    # get next step actionddd
-                    objects = self.update_objects(objects, arm=arm)
+                        set_gripper_open = gripper_action == 'open'
+                        set_gripper_close = gripper_action == 'close'
+                        self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
+                                                          self.attached_obj_id is not None, set_gripper_open,
+                                                          set_gripper_close, arm=arm, target_pose=target_gripper_pose)
+                        if last_gripper_state != gripper_action and gripper_action is not None:
+                            self.robot.set_gripper_action(gripper_action, arm=arm)
+                            last_gripper_state = gripper_action
+                            if gripper_action == 'open':
+                                time.sleep(1)
 
-                    target_gripper_pose, motion_type, gripper_action, arm = substages.get_action(objects)
-                    arm = extra_params.get('arm', 'right')
-                    self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
-                                                      self.attached_obj_id is not None, arm=arm,
-                                                      target_pose=target_gripper_pose)
+                            self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
+                                                            self.attached_obj_id is not None, arm=arm,
+                                                            target_pose=target_gripper_pose)
 
-                    # execution action
-                    if target_gripper_pose is not None:
-                        self.robot.move_pose(target_gripper_pose, motion_type, arm=arm, block=True)
+                        # check sub-stage completion
+                        objects['gripper'].obj_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
+                        objects = self.update_objects(objects, arm=arm)
 
-                    set_gripper_open = gripper_action == 'open'
-                    set_gripper_close = gripper_action == 'close'
-                    self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
-                                                      self.attached_obj_id is not None, set_gripper_open,
-                                                      set_gripper_close, arm=arm, target_pose=target_gripper_pose)
-                    self.robot.set_gripper_action(gripper_action, arm=arm)
-                    if gripper_action == 'open':
-                        time.sleep(1)
+                        success = substages.check_completion(objects)
+                        self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
+                                                          self.attached_obj_id is not None, arm=arm,
+                                                          target_pose=target_gripper_pose)
+                        if success == False:
+                            # import ipdb;ipdb.set_trace()
+                            self.attached_obj_id = None
+                            Log.warning('Failed at sub-stage %d' % substages.step_id)
+                            break
 
-                    self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
-                                                      self.attached_obj_id is not None, arm=arm,
-                                                      target_pose=target_gripper_pose)
+                        # attach grasped object to gripper           # TODO avoid articulated objects
+                        if gripper_action == 'close':  # TODO  确定是grasp才行！！
+                            self.attached_obj_id = substages.passive_obj_id
+                        elif gripper_action == 'open':
+                            self.attached_obj_id = None
+                        self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
+                                                          self.attached_obj_id is not None, arm=arm,
+                                                          target_pose=target_gripper_pose)
 
-                    # check sub-stage completion
-                    objects['gripper'].obj_pose = self.robot.get_ee_pose(ee_type='gripper', id=arm)
-                    objects = self.update_objects(objects, arm=arm)
-
-                    success = substages.check_completion(objects)
-                    self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
-                                                      self.attached_obj_id is not None, arm=arm,
-                                                      target_pose=target_gripper_pose)
-                    if success == False:
-                        # import ipdb;ipdb.set_trace()
-                        self.attached_obj_id = None
-                        print('Failed at sub-stage %d' % substages.step_id)
-                        break
-
-                    # attach grasped object to gripper           # TODO avoid articulated objects
-                    if gripper_action == 'close':  # TODO  确定是grasp才行！！
-                        self.attached_obj_id = substages.passive_obj_id
-                    elif gripper_action == 'open':
-                        self.attached_obj_id = None
-                    self.robot.client.set_frame_state(action, substages.step_id, active_id, passive_id,
-                                                      self.attached_obj_id is not None, arm=arm,
-                                                      target_pose=target_gripper_pose)
-
-                    # change object position
-                    num_against = substages.extra_params.get('against', 0)
-                    against_area = substages.extra_params.get('against_range', [])
-                    against_type = substages.extra_params.get('against_type', None)
-                    if num_against > 0 and gripper_action == 'open' and action == 'pick' and (
-                            against_type is not None):
-                        parts = against_type.split('_')
-                        need_move_objects = [passive_id]
-                        if parts[0] == 'move' and parts[1] == 'with':
-                            for workspace in objects:
-                                if parts[2] in workspace:
-                                    need_move_objects.append(workspace)
-                        ## 目前就集体向一个方向移动
-                        offset_y = random.uniform(0, 0.2)
-                        poses = []
-                        for move_object in need_move_objects:
-                            response = self.robot.client.get_object_pose(f'/World/Objects/{move_object}')
+                        # change object position
+                        num_against = substages.extra_params.get('against', 0)
+                        against_area = substages.extra_params.get('against_range', [])
+                        against_type = substages.extra_params.get('against_type', None)
+                        if num_against > 0 and gripper_action == 'open' and action == 'pick' and (
+                                against_type is not None):
+                            parts = against_type.split('_')
+                            need_move_objects = [passive_id]
+                            if parts[0] == 'move' and parts[1] == 'with':
+                                for workspace in objects:
+                                    if parts[2] in workspace:
+                                        need_move_objects.append(workspace)
+                            ## 目前就集体向一个方向移动
+                            offset_y = random.uniform(0, 0.2)
+                            poses = []
+                            for move_object in need_move_objects:
+                                response = self.robot.client.get_object_pose(f'/World/Objects/{move_object}')
+                                pos = [response.object_pose.position.x,
+                                       response.object_pose.position.y + offset_y,
+                                       response.object_pose.position.z]
+                                quat_wxyz = np.array(
+                                    [
+                                        response.object_pose.rpy.rw,
+                                        response.object_pose.rpy.rx,
+                                        response.object_pose.rpy.ry,
+                                        response.object_pose.rpy.rz,
+                                    ]
+                                )
+                                object_pose = {}
+                                object_pose["prim_path"] = f'/World/Objects/{move_object}'
+                                object_pose["position"] = pos
+                                object_pose["rotation"] = quat_wxyz
+                                poses.append(object_pose)
+                            print(poses)
+                            self.robot.client.SetObjectPose(poses, [])
+                            self.robot.client.DetachObj()
+                        elif (num_against > 0 and gripper_action == 'open' and action == 'pick') or \
+                                (num_against > 0 and action == 'place' and gripper_action == 'close'):
+                            # import ipdb;ipdb.set_trace()
+                            x_min, x_max = 999.0, -999.0
+                            y_min, y_max = 999.0, -999.0
+                            if against_area:
+                                selected_against_area = random.choice(against_area)
+                                if selected_against_area in workspaces:
+                                    position = workspaces[selected_against_area]['position']
+                                    size = workspaces[selected_against_area]['size']
+                                    x_min, x_max = position[0] - size[0] / 2, position[0] + size[0] / 2
+                                    y_min, y_max = position[1] - size[1] / 2, position[1] + size[1] / 2
+                                    x_size, y_size, z_size = objects[passive_id].info['size']
+                                    up_axis = objects[passive_id].info['upAxis']
+                                    axis_mapping = {
+                                        'x': (y_size / 2000.0, z_size / 2000.0),
+                                        'y': (x_size / 2000.0, z_size / 2000.0),
+                                        'z': (x_size / 2000.0, y_size / 2000.0)
+                                    }
+                                    dimensions = axis_mapping[up_axis[0]]
+                                    distance = np.linalg.norm(dimensions)
+                                    x_min += distance * 1.5
+                                    x_max -= distance * 1.5
+                                    y_min += distance * 1.5
+                                    y_max -= distance * 1.5
+                            else:
+                                print("against_range not set")
+                                continue
+                            response = self.robot.client.get_object_pose(f'/World/Objects/{passive_id}')
                             pos = [response.object_pose.position.x,
-                                   response.object_pose.position.y + offset_y,
-                                   response.object_pose.position.z]
+                                   response.object_pose.position.y,
+                                   response.object_pose.position.z + 0.02]
                             quat_wxyz = np.array(
                                 [
                                     response.object_pose.rpy.rw,
@@ -339,121 +423,44 @@ class Agent(BaseAgent):
                                     response.object_pose.rpy.rz,
                                 ]
                             )
+                            # import ipdb;ipdb.set_trace()
+                            # if position is close to the gripper, then random position again
+                            while True:
+                                pos[0] = random.uniform(x_min, x_max)
+                                pos[1] = random.uniform(y_min, y_max)
+                                distance = np.linalg.norm(
+                                    [pos[0] - target_gripper_pose[0][3], pos[1] - target_gripper_pose[1][3]])
+                                if distance >= 0.2:
+                                    break
+                            poses = []
                             object_pose = {}
-                            object_pose["prim_path"] = f'/World/Objects/{move_object}'
+                            object_pose["prim_path"] = f'/World/Objects/{passive_id}'
                             object_pose["position"] = pos
                             object_pose["rotation"] = quat_wxyz
                             poses.append(object_pose)
-                        print(poses)
-                        self.robot.client.SetObjectPose(poses, [])
-                        self.robot.client.DetachObj()
-                    elif (num_against > 0 and gripper_action == 'open' and action == 'pick') or \
-                            (num_against > 0 and action == 'place' and gripper_action == 'close'):
-                        # import ipdb;ipdb.set_trace()
-                        x_min, x_max = 999.0, -999.0
-                        y_min, y_max = 999.0, -999.0
-                        if against_area:
-                            selected_against_area = random.choice(against_area)
-                            if selected_against_area in workspaces:
-                                position = workspaces[selected_against_area]['position']
-                                size = workspaces[selected_against_area]['size']
-                                x_min, x_max = position[0] - size[0] / 2, position[0] + size[0] / 2
-                                y_min, y_max = position[1] - size[1] / 2, position[1] + size[1] / 2
-                                x_size, y_size, z_size = objects[passive_id].info['size']
-                                up_axis = objects[passive_id].info['upAxis']
-                                axis_mapping = {
-                                    'x': (y_size / 2000.0, z_size / 2000.0),
-                                    'y': (x_size / 2000.0, z_size / 2000.0),
-                                    'z': (x_size / 2000.0, y_size / 2000.0)
-                                }
-                                dimensions = axis_mapping[up_axis[0]]
-                                distance = np.linalg.norm(dimensions)
-                                x_min += distance * 1.5
-                                x_max -= distance * 1.5
-                                y_min += distance * 1.5
-                                y_max -= distance * 1.5
-                        else:
-                            print("against_range not set")
-                            continue
-                        response = self.robot.client.get_object_pose(f'/World/Objects/{passive_id}')
-                        pos = [response.object_pose.position.x,
-                               response.object_pose.position.y,
-                               response.object_pose.position.z + 0.02]
-                        quat_wxyz = np.array(
-                            [
-                                response.object_pose.rpy.rw,
-                                response.object_pose.rpy.rx,
-                                response.object_pose.rpy.ry,
-                                response.object_pose.rpy.rz,
-                            ]
-                        )
-                        # import ipdb;ipdb.set_trace()
-                        # if position is close to the gripper, then random position again
-                        while True:
-                            pos[0] = random.uniform(x_min, x_max)
-                            pos[1] = random.uniform(y_min, y_max)
-                            distance = np.linalg.norm(
-                                [pos[0] - target_gripper_pose[0][3], pos[1] - target_gripper_pose[1][3]])
-                            if distance >= 0.2:
-                                break
-                        poses = []
-                        object_pose = {}
-                        object_pose["prim_path"] = f'/World/Objects/{passive_id}'
-                        object_pose["position"] = pos
-                        object_pose["rotation"] = quat_wxyz
-                        poses.append(object_pose)
-                        print(poses)
-                        self.robot.client.SetObjectPose(poses, [])
-                        self.robot.client.DetachObj()
-                    # if num_against > 0 and action == 'place' and gripper_action == 'close':
-                    #     offset = random.uniform(-0.1, 0.1)
-                    #     response = self.robot.client.get_object_pose('/World/Objects/%s'%passive_id)
-                    #     pos = [response.object_pose.position.x,
-                    #            response.object_pose.position.y - offset,
-                    #            response.object_pose.position.z]
-                    #     quat_wxyz = np.array(
-                    #         [
-                    #             response.object_pose.rpy.rw,
-                    #             response.object_pose.rpy.rx,
-                    #             response.object_pose.rpy.ry,
-                    #             response.object_pose.rpy.rz,
-                    #         ]
-                    #     )
-                    #     # import ipdb;ipdb.set_trace()
-                    #     poses = []
-                    #     object_pose ={}
-                    #     object_pose["prim_path"] = '/World/Objects/' + passive_id
-                    #     object_pose["position"] =pos
-                    #     object_pose["rotation"] = quat_wxyz
-                    #     poses.append(object_pose)
-                    #     self.robot.client.SetObjectPose(poses, [])
-                    #     index += 1
-                    #     self.robot.client.DetachObj()
-
+                            print(poses)
+                            self.robot.client.SetObjectPose(poses, [])
+                            self.robot.client.DetachObj()
+                    if success == False:
+                        self.attached_obj_id = None
+                        break
+                    if self.attached_obj_id is not None:
+                        if self.attached_obj_id.split('/')[0] not in self.articulated_objs:
+                            self.robot.client.DetachObj()
+                            self.robot.client.AttachObj(prim_paths=['/World/Objects/' + self.attached_obj_id],
+                                                        is_right=arm == 'right')
                 if success == False:
-                    self.attached_obj_id = None
                     break
-                if self.attached_obj_id is not None:
-                    if self.attached_obj_id.split('/')[0] not in self.articulated_objs:
-                        self.robot.client.DetachObj()
-                        self.robot.client.AttachObj(prim_paths=['/World/Objects/' + self.attached_obj_id],
-                                                    is_right=arm == 'right')
-            if success == False:
-                break
-        time.sleep(0.5)
-        if self.attached_obj_id is None:
-            self.robot.client.DetachObj()
-        self.robot.client.stop_recording()
-        # try:
-        #     step_id = substages.step_id if substages is not None and len(substages) else -1
-        # except:
-        #     step_id = -1
-        step_id = -1
-        fail_stage_step = [stage_id, step_id] if success == False else [-1, -1]
+            Log.success('Finished executing task.')
+            time.sleep(0.5)
+            #import ipdb; ipdb.set_trace()
+            if self.attached_obj_id is None:
+                self.robot.client.DetachObj()
+            self.robot.client.stop_recording()
+            step_id = -1
+            fail_stage_step = [stage_id, step_id] if not success else [-1, -1]
 
-        task_info_saved = task_info.copy()
-        self.robot.client.SendTaskStatus(success, fail_stage_step)
-        if success:
-            print(">>>>>>>>>>>>>>>> Success!!!!!!!!!!!!!!!!")
+            task_info_saved = task_info.copy()
+            self.robot.client.SendTaskStatus(success, fail_stage_step)
 
         return True

data_gen_dependencies/robot_urdf/Franka_120s.urdf

@@ -0,0 +1,877 @@
+<?xml version="1.0" ?>
+<robot name="panda" xmlns:xacro="http://www.ros.org/wiki/xacro">
+  <link name="base_link"/>
+  <joint name="panda_fixed" type="fixed">
+    <origin rpy="0 0 0.0" xyz="0 0 0.0"/>
+    <parent link="base_link"/>
+    <child link="panda_link0"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+
+  <link name="panda_link0">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link0.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link0.obj"/>
+      </geometry>
+    </collision>
+
+
+  </link>
+  <link name="panda_link1">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link1.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link1.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.00033 -0.02204 -0.04762"/>
+      <mass value="2.797"/>
+      <inertia ixx="0.0156" ixy="0.0" ixz="0.0" iyy="0.01439" iyz="-0.0024" izz="0.005"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint1" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="0 0 0" xyz="0 0 0.333"/>
+    <parent link="panda_link0"/>
+    <child link="panda_link1"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-2.8973" upper="2.8973" velocity="2.1750"/>
+  </joint>
+  <link name="panda_link2">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link2.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link2.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.00038 -0.09211 0.01908"/>
+      <mass value="2.542"/>
+      <inertia ixx="0.0166" ixy="0.0" ixz="0.0" iyy="0.0046" iyz="0.0035" izz="0.015"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint2" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-1.7628" soft_upper_limit="1.7628"/>
+    <origin rpy="-1.57079632679 0 0" xyz="0 0 0"/>
+    <parent link="panda_link1"/>
+    <child link="panda_link2"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-1.7628" upper="1.7628" velocity="2.1750"/>
+  </joint>
+  <link name="panda_link3">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link3.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link3.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.05152 0.01696 -0.02971"/>
+      <mass value="2.2513"/>
+      <inertia ixx="0.006" ixy="0.009" ixz="0.003" iyy="0.0086" iyz="0.009" izz="0.0065"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint3" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="1.57079632679 0 0" xyz="0 -0.316 0"/>
+    <parent link="panda_link2"/>
+    <child link="panda_link3"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-2.8973" upper="2.8973" velocity="2.1750"/>
+  </joint>
+  <link name="panda_link4">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link4.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link4.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.05113 0.05825 0.01698"/>
+      <mass value="2.2037"/>
+      <inertia ixx="0.0156" ixy="0.0" ixz="0.0" iyy="0.01439" iyz="-0.0024" izz="0.005"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint4" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-3.0718" soft_upper_limit="-0.0698"/>
+    <origin rpy="1.57079632679 0 0" xyz="0.0825 0 0"/>
+    <parent link="panda_link3"/>
+    <child link="panda_link4"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-3.0718" upper="-0.0698" velocity="2.1750"/>
+    <!-- something is weird with this joint limit config
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-3.0" upper="0.087" velocity="2.1750"/>  -->
+  </joint>
+  <link name="panda_link5">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link5.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link5.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.00005 0.03730 -0.09280"/>
+      <mass value="2.2855"/>
+      <inertia ixx="0.02297014781" ixy="0.0" ixz="0.0" iyy="0.02095060919" iyz="0.00382345782" izz="0.00430606551"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint5" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="-1.57079632679 0 0" xyz="-0.0825 0.384 0"/>
+    <parent link="panda_link4"/>
+    <child link="panda_link5"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="12" lower="-2.8973" upper="2.8973" velocity="2.6100"/>
+  </joint>
+  <link name="panda_link6">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link6.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link6.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.06572 -0.00371 0.00153"/>
+      <mass value="1.353"/>
+      <inertia ixx="0.00087964522" ixy="-0.00021487814" ixz="-0.00011911662" iyy="0.00277796968" iyz="0.00001274322" izz="0.00286701969"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint6" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-0.0175" soft_upper_limit="3.7525"/>
+    <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
+    <parent link="panda_link5"/>
+    <child link="panda_link6"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="12" lower="-0.0175" upper="3.7525" velocity="2.6100"/>
+    <!-- <dynamics damping="10.0"/>
+    <limit effort="12" lower="-0.0873" upper="3.0" velocity="2.6100"/> -->
+  </joint>
+  <link name="panda_link7">
+    <visual>
+      <geometry>
+        <mesh filename="meshes/visual/link7.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="meshes/collision/link7.obj"/>
+      </geometry>
+    </collision>
+        <inertial>
+      <origin rpy="0 0 0" xyz="0.00089 -0.00044 0.05491"/>
+      <mass value="0.35973"/>
+      <inertia ixx="0.00019541063" ixy="0.00000165231" ixz="0.00000148826" iyy="0.00019210361" iyz="-0.00000131132" izz="0.00017936256"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint7" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="1.57079632679 0 0" xyz="0.088 0 0"/>
+    <parent link="panda_link6"/>
+    <child link="panda_link7"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="12" lower="-2.8973" upper="2.8973" velocity="2.6100"/>
+  </joint>
+
+  <link name="panda_link8"/>
+  <joint name="panda_joint8" type="fixed">
+    <origin rpy="0 0 0" xyz="0 0 0.107"/>
+    <parent link="panda_link7"/>
+    <child link="panda_link8"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  <joint name="panda_hand_joint" type="fixed">
+    <parent link="panda_link8"/>
+    <child link="right_base_link"/>
+
+    <origin rpy="0 0 0.785398163397" xyz="0 0 0"/>
+  </joint>
+
+  <link
+      name="right_base_link">
+      <inertial>
+        <origin
+          xyz="1.23989288637867E-06 3.27564415436087E-07 0.0536186975993762"
+          rpy="0 0 0" />
+        <mass
+          value="0.443105756928829" />
+        <inertia
+          ixx="0.00064200574020552"
+          ixy="-1.51180224931458E-07"
+          ixz="7.24174172220287E-09"
+          iyy="0.00066998312004943"
+          iyz="2.73584170951165E-09"
+          izz="0.000278949393571484" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/base_link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/base_link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <link
+      name="right_Left_01_Link">
+      <inertial>
+        <origin
+          xyz="0.00656387870913618 0.024850259757102 -9.40817030196858E-10"
+          rpy="0 0 0" />
+        <mass
+          value="0.0112899386472626" />
+        <inertia
+          ixx="1.96944973967669E-06"
+          ixy="6.42395813223716E-07"
+          ixz="-1.07950451056367E-14"
+          iyy="1.39704780795747E-06"
+          iyz="-1.48209813327181E-14"
+          izz="2.47457179553145E-06" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_01_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_01_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Left_1_Joint"
+      type="revolute">
+      <origin
+        xyz="-0.032553 0 0.10644"
+        rpy="1.5708 0 3.1416" />
+      <parent
+        link="right_base_link" />
+      <child
+        link="right_Left_01_Link" />
+      <axis
+        xyz="0 0 -1" />
+      <limit
+        lower="0"
+        upper="1"
+        effort="2"
+        velocity="1000" />
+    </joint>
+
+
+    <link
+      name="right_Left_00_Link">
+      <inertial>
+        <origin
+          xyz="0.00656387870913618 0.024850259757102 -9.40817030196858E-10"
+          rpy="0 0 0" />
+        <mass
+          value="0.0112899386472626" />
+        <inertia
+          ixx="1.96944973967669E-06"
+          ixy="6.42395813223716E-07"
+          ixz="-1.07950451056367E-14"
+          iyy="1.39704780795747E-06"
+          iyz="-1.48209813327181E-14"
+          izz="2.47457179553145E-06" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_00_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_00_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+
+    <joint
+      name="right_Left_0_Joint"
+      type="revolute">
+      <origin
+        xyz="0.01946 0.0129 0"
+        rpy="0 0 0" />
+      <parent
+        link="right_Left_01_Link" />
+      <child
+        link="right_Left_00_Link" />
+      <axis
+        xyz="0 0 -1" />
+      <limit
+        lower="0"
+        upper="1"
+        effort="2"
+        velocity="1000" />
+    </joint>
+
+
+    <link
+      name="right_Left_Support_Link">
+      <inertial>
+        <origin
+          xyz="-0.011970095332009 0.0238338043684485 -1.36624428992904E-06"
+          rpy="0 0 0" />
+        <mass
+          value="0.0140092102116716" />
+        <inertia
+          ixx="4.13674978137002E-06"
+          ixy="9.81573398608352E-07"
+          ixz="-2.87921970465553E-11"
+          iyy="1.1362453186293E-06"
+          iyz="2.94726072168702E-11"
+          izz="4.10226218364235E-06" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_Support_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_Support_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Left_Support_Joint"
+      type="revolute">
+      <origin
+        xyz="-0.02888 0.04341 0"
+        rpy="0 0 0" />
+      <parent
+        link="right_Left_00_Link" />
+      <child
+        link="right_Left_Support_Link" />
+      <axis
+        xyz="0 0 1" />
+      <limit
+        lower="-1"
+        upper="1"
+        effort="10"
+        velocity="10" />
+    </joint>
+    <link
+      name="right_Left_Pad_Link">
+      <inertial>
+        <origin
+          xyz="1.46688816627348E-09 0.00209017783344978 -8.17280129359337E-17"
+          rpy="0 0 0" />
+        <mass
+          value="0.00332411394949541" />
+        <inertia
+          ixx="5.5867126170785E-07"
+          ixy="2.9966962370347E-20"
+          ixz="1.44978830317617E-23"
+          iyy="1.45407135182096E-07"
+          iyz="-1.60867378777053E-21"
+          izz="4.18250297449997E-07" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_Pad_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.384313725490196 0.384313725490196 0.384313725490196 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_Pad_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Left_Pad_Joint"
+      type="fixed">
+      <origin
+        xyz="-0.022056 0.044306 0"
+        rpy="0 0 0" />
+      <parent
+        link="right_Left_Support_Link" />
+      <child
+        link="right_Left_Pad_Link" />
+      <axis
+        xyz="0 0 0" />
+    </joint>
+    <link
+      name="right_Left_2_Link">
+      <inertial>
+        <origin
+          xyz="-0.00438952959966343 0.0248815722409523 -2.55135883071891E-05"
+          rpy="0 0 0" />
+        <mass
+          value="0.00803955821311802" />
+        <inertia
+          ixx="1.15079028835042E-06"
+          ixy="1.13853775681911E-07"
+          ixz="-1.79492796413463E-10"
+          iyy="4.68655200337902E-07"
+          iyz="-4.49903912976825E-10"
+          izz="8.47367915692972E-07" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_2_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Left_2_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Left_2_Joint"
+      type="revolute">
+      <origin
+        xyz="-0.017 0 0.122"
+        rpy="1.5708 0 3.1416" />
+      <parent
+        link="right_base_link" />
+      <child
+        link="right_Left_2_Link" />
+      <axis
+        xyz="0 0 -1" />
+      <limit
+        lower="-1"
+        upper="1"
+        effort="10"
+        velocity="10" />
+    </joint>
+    <link
+      name="right_Right_2_Link">
+      <inertial>
+        <origin
+          xyz="-0.00438952959968201 0.0248815722409522 -2.5513588307359E-05"
+          rpy="0 0 0" />
+        <mass
+          value="0.00803955821311803" />
+        <inertia
+          ixx="1.15079028835042E-06"
+          ixy="1.13853775681912E-07"
+          ixz="-1.79492796413511E-10"
+          iyy="4.68655200337903E-07"
+          iyz="-4.49903912976901E-10"
+          izz="8.47367915692974E-07" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_2_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_2_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Right_2_Joint"
+      type="revolute">
+      <origin
+        xyz="0.017 0 0.122"
+        rpy="1.5708 0 0" />
+      <parent
+        link="right_base_link" />
+      <child
+        link="right_Right_2_Link" />
+      <axis
+        xyz="0 0 1" />
+      <limit
+        lower="-1"
+        upper="1"
+        effort="10"
+        velocity="10" />
+    </joint>
+    <link
+      name="right_Right_01_Link">
+      <inertial>
+        <origin
+          xyz="0.00656387870564603 0.0248502597572661 -9.40812230336897E-10"
+          rpy="0 0 0" />
+        <mass
+          value="0.0112899386472626" />
+        <inertia
+          ixx="1.96944973967663E-06"
+          ixy="6.42395813223747E-07"
+          ixz="-1.07950449101923E-14"
+          iyy="1.39704780795753E-06"
+          iyz="-1.48209815851649E-14"
+          izz="2.47457179553146E-06" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_01_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_01_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Right_1_Joint"
+      type="revolute">
+      <origin
+        xyz="0.032553 0 0.10644"
+        rpy="1.5708 0 0" />
+      <parent
+        link="right_base_link" />
+      <child
+        link="right_Right_01_Link" />
+      <axis
+        xyz="0 0 1" />
+      <limit
+        lower="-1"
+        upper="1"
+        effort="10"
+        velocity="10" />
+    </joint>
+
+    <link
+      name="right_Right_00_Link">
+      <inertial>
+        <origin
+          xyz="0.00656387870564603 0.0248502597572661 -9.40812230336897E-10"
+          rpy="0 0 0" />
+        <mass
+          value="0.0112899386472626" />
+        <inertia
+          ixx="1.96944973967663E-06"
+          ixy="6.42395813223747E-07"
+          ixz="-1.07950449101923E-14"
+          iyy="1.39704780795753E-06"
+          iyz="-1.48209815851649E-14"
+          izz="2.47457179553146E-06" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_00_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_00_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+
+     <joint
+      name="right_Right_0_Joint"
+      type="revolute">
+      
+      <origin
+        xyz="0.01946 0.0129 0"
+        rpy="0 0 0" />
+      <parent
+        link="right_Right_01_Link" />
+      <child
+        link="right_Right_00_Link" />
+      <axis
+        xyz="0 0 1" />
+      <limit
+        lower="-1"
+        upper="1"
+        effort="10"
+        velocity="10" />
+    </joint>
+
+    <link
+      name="right_Right_Support_Link">
+      <inertial>
+        <origin
+          xyz="-0.011970095112324 0.0238338043682057 -1.36624428997256E-06"
+          rpy="0 0 0" />
+        <mass
+          value="0.0140092102116716" />
+        <inertia
+          ixx="4.13674978136988E-06"
+          ixy="9.81573398608576E-07"
+          ixz="-2.87921970465403E-11"
+          iyy="1.13624531862944E-06"
+          iyz="2.94726072165632E-11"
+          izz="4.10226218364236E-06" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_Support_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.792156862745098 0.819607843137255 0.933333333333333 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_Support_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Right_Support_Joint"
+      type="revolute">
+      <origin
+        xyz="-0.02888 0.04341 0"
+        rpy="0 0 0" />
+      <parent
+        link="right_Right_00_Link" />
+      <child
+        link="right_Right_Support_Link" />
+      <axis
+        xyz="0 0 1" />
+      <limit
+        lower="-1"
+        upper="1"
+        effort="10"
+        velocity="10" />
+    </joint>
+    <link
+      name="right_Right_Pad_Link">
+      <inertial>
+        <origin
+          xyz="5.43457184440897E-09 0.00209017783344967 -1.21646039264426E-16"
+          rpy="0 0 0" />
+        <mass
+          value="0.00332411394949542" />
+        <inertia
+          ixx="5.58671261707851E-07"
+          ixy="3.06066210315894E-20"
+          ixz="1.70379713307979E-23"
+          iyy="1.45407135182096E-07"
+          iyz="-1.65096420438087E-21"
+          izz="4.18250297449998E-07" />
+      </inertial>
+      <visual>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_Pad_Link.STL" />
+        </geometry>
+        <material
+          name="">
+          <color
+            rgba="0.384313725490196 0.384313725490196 0.384313725490196 1" />
+        </material>
+      </visual>
+      <collision>
+        <origin
+          xyz="0 0 0"
+          rpy="0 0 0" />
+        <geometry>
+          <mesh
+            filename="package://meshes/crt_120s/Right_Pad_Link.STL" />
+        </geometry>
+      </collision>
+    </link>
+    <joint
+      name="right_Right_Pad_Joint"
+      type="fixed">
+      <origin
+        xyz="-0.022056 0.044306 0"
+        rpy="0 0 0" />
+      <parent
+        link="right_Right_Support_Link" />
+      <child
+        link="right_Right_Pad_Link" />
+      <axis
+        xyz="0 0 0" />
+    </joint>
+
+
+
+
+  <link name="ee_link"/>
+
+  <joint name="ee_fixed_joint" type="fixed">
+    <parent link="right_base_link"/>
+    <child link="ee_link"/>
+    <origin rpy="0 0.0 -1.57079632679" xyz="0.0 0. 0.225"/>
+  </joint>
+  <link name="right_gripper">
+    <inertial>
+      <!-- Dummy inertial parameters to avoid link lumping-->
+      <mass value="0.01"/>
+      <inertia ixx="0.0001" ixy="0.0" ixz="0.0" iyy="0.0001" iyz="0.0" izz="0.0001"/>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+    </inertial>
+  </link>
+  <joint name="right_gripper" type="fixed">
+    <!--<origin rpy="0 0 2.35619449019" xyz="0 0 0.1"/>-->
+    <origin rpy="0 0 2.35619449019" xyz="0 0 0.207"/>
+    <axis xyz="0 0 1"/>
+    <parent link="panda_link7"/>
+    <!--<parent link="panda_link8"/>-->
+    <child link="right_gripper"/>
+  </joint>
+
+
+</robot>
+
+

data_gen_dependencies/robot_urdf/Franka_120s_robotiq_2f85.urdf

@@ -0,0 +1,531 @@
+<?xml version="1.0" ?>
+<robot name="panda" xmlns:xacro="http://www.ros.org/wiki/xacro">
+  <link name="base_link"/>
+  <joint name="panda_fixed" type="fixed">
+    <origin rpy="0 0 0.0" xyz="0 0 0.0"/>
+    <parent link="base_link"/>
+    <child link="panda_link0"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+
+  <link name="panda_link0">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link0.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link0.obj"/>
+      </geometry>
+    </collision>
+
+
+  </link>
+  <link name="panda_link1">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link1.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link1.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.00033 -0.02204 -0.04762"/>
+      <mass value="2.797"/>
+      <inertia ixx="0.0156" ixy="0.0" ixz="0.0" iyy="0.01439" iyz="-0.0024" izz="0.005"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint1" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="0 0 0" xyz="0 0 0.333"/>
+    <parent link="panda_link0"/>
+    <child link="panda_link1"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-2.8973" upper="2.8973" velocity="2.1750"/>
+  </joint>
+  <link name="panda_link2">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link2.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link2.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.00038 -0.09211 0.01908"/>
+      <mass value="2.542"/>
+      <inertia ixx="0.0166" ixy="0.0" ixz="0.0" iyy="0.0046" iyz="0.0035" izz="0.015"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint2" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-1.7628" soft_upper_limit="1.7628"/>
+    <origin rpy="-1.57079632679 0 0" xyz="0 0 0"/>
+    <parent link="panda_link1"/>
+    <child link="panda_link2"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-1.7628" upper="1.7628" velocity="2.1750"/>
+  </joint>
+  <link name="panda_link3">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link3.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link3.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.05152 0.01696 -0.02971"/>
+      <mass value="2.2513"/>
+      <inertia ixx="0.006" ixy="0.009" ixz="0.003" iyy="0.0086" iyz="0.009" izz="0.0065"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint3" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="1.57079632679 0 0" xyz="0 -0.316 0"/>
+    <parent link="panda_link2"/>
+    <child link="panda_link3"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-2.8973" upper="2.8973" velocity="2.1750"/>
+  </joint>
+  <link name="panda_link4">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link4.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link4.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.05113 0.05825 0.01698"/>
+      <mass value="2.2037"/>
+      <inertia ixx="0.0156" ixy="0.0" ixz="0.0" iyy="0.01439" iyz="-0.0024" izz="0.005"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint4" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-3.0718" soft_upper_limit="-0.0698"/>
+    <origin rpy="1.57079632679 0 0" xyz="0.0825 0 0"/>
+    <parent link="panda_link3"/>
+    <child link="panda_link4"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-3.0718" upper="-0.0698" velocity="2.1750"/>
+    <!-- something is weird with this joint limit config
+    <dynamics damping="10.0"/>
+    <limit effort="87" lower="-3.0" upper="0.087" velocity="2.1750"/>  -->
+  </joint>
+  <link name="panda_link5">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link5.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link5.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.00005 0.03730 -0.09280"/>
+      <mass value="2.2855"/>
+      <inertia ixx="0.02297014781" ixy="0.0" ixz="0.0" iyy="0.02095060919" iyz="0.00382345782" izz="0.00430606551"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint5" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="-1.57079632679 0 0" xyz="-0.0825 0.384 0"/>
+    <parent link="panda_link4"/>
+    <child link="panda_link5"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="12" lower="-2.8973" upper="2.8973" velocity="2.6100"/>
+  </joint>
+  <link name="panda_link6">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link6.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link6.obj"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.06572 -0.00371 0.00153"/>
+      <mass value="1.353"/>
+      <inertia ixx="0.00087964522" ixy="-0.00021487814" ixz="-0.00011911662" iyy="0.00277796968" iyz="0.00001274322" izz="0.00286701969"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint6" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-0.0175" soft_upper_limit="3.7525"/>
+    <origin rpy="1.57079632679 0 0" xyz="0 0 0"/>
+    <parent link="panda_link5"/>
+    <child link="panda_link6"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="12" lower="-0.0175" upper="3.7525" velocity="2.6100"/>
+    <!-- <dynamics damping="10.0"/>
+    <limit effort="12" lower="-0.0873" upper="3.0" velocity="2.6100"/> -->
+  </joint>
+  <link name="panda_link7">
+    <visual>
+      <geometry>
+        <mesh filename="franka_meshes/visual/link7.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="franka_meshes/collision/link7.obj"/>
+      </geometry>
+    </collision>
+        <inertial>
+      <origin rpy="0 0 0" xyz="0.00089 -0.00044 0.05491"/>
+      <mass value="0.35973"/>
+      <inertia ixx="0.00019541063" ixy="0.00000165231" ixz="0.00000148826" iyy="0.00019210361" iyz="-0.00000131132" izz="0.00017936256"/>
+    </inertial>
+  </link>
+  <joint name="panda_joint7" type="revolute">
+    <safety_controller k_position="100.0" k_velocity="40.0" soft_lower_limit="-2.8973" soft_upper_limit="2.8973"/>
+    <origin rpy="1.57079632679 0 0" xyz="0.088 0 0"/>
+    <parent link="panda_link6"/>
+    <child link="panda_link7"/>
+    <axis xyz="0 0 1"/>
+    <dynamics damping="10.0"/>
+    <limit effort="12" lower="-2.8973" upper="2.8973" velocity="2.6100"/>
+  </joint>
+
+  <link name="panda_link8"/>
+  <joint name="panda_joint8" type="fixed">
+    <origin rpy="0 0 0" xyz="0 0 0.107"/>
+    <parent link="panda_link7"/>
+    <child link="panda_link8"/>
+    <axis xyz="0 0 0"/>
+  </joint>
+  
+   <link name="flange">
+    <visual>
+      <geometry>
+        <mesh filename="flange_meshes/flange.obj"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="flange_meshes/flange.obj"/>
+      </geometry>
+    </collision>
+        <inertial>
+      <origin rpy="0 0 0" xyz="0.0 0 0.0"/>
+      <mass value="0.1"/>
+      <inertia ixx="0.0001" ixy="0" ixz="0" iyy="0.0001" iyz="0" izz="0.0001"/>
+    </inertial>
+  </link>
+  
+  <link name="zed_mini_camera">
+    <visual>
+      <geometry>
+        <mesh filename="zed_mini_camera_meshes/zed_mini_camera.obj" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="zed_mini_camera_meshes/zed_mini_camera.obj" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+        <inertial>
+      <origin rpy="0 0 0" xyz="0.0 0 0.0"/>
+      <mass value="0.1"/>
+      <inertia ixx="0.0001" ixy="0" ixz="0" iyy="0.0001" iyz="0" izz="0.0001"/>
+    </inertial>
+  </link>
+
+  <joint name="panda_zed_mini_camera_joint" type="fixed">
+    <parent link="panda_link8"/>
+    <child link="zed_mini_camera"/>
+    <origin xyz="0 0 -0.01" rpy="-1.57079632679 -1.57079632679 0"/>
+  </joint>
+
+  <!-- flange 连接在 link8 上 -->
+  <joint name="panda_flange_joint" type="fixed">
+    <origin rpy="-1.57079632679 0 0" xyz="0 0 0.003"/>
+    <parent link="panda_link8"/>
+    <child link="flange"/>
+  </joint>
+
+<!-- gripper 连接在 flange 上 -->
+  <joint name="flange_robotiq_fix_joint" type="fixed">
+    <origin rpy="1.57079632679 1.57079632679 0" xyz="0 -0.0035 0"/>
+    <parent link="flange"/>
+    <child link="robotiq_85_base_link"/>
+  </joint>
+
+  
+  <link name="robotiq_85_base_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/robotiq_base.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/robotiq_base.stl"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.0 2.274e-05 0.03232288"/>
+      <mass value="6.6320197e-01"/>
+      <inertia ixx="5.1617816e-04" ixy="2.936e-8" ixz="0.0" iyy="5.8802208e-04" iyz="-3.2296e-7" izz="3.9462776e-04"/>
+    </inertial>
+  </link>
+
+  <link name="robotiq_85_left_knuckle_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/left_knuckle.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/left_knuckle.stl"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.01213197 0.0002 -0.00058647"/>
+      <mass value="1.384773208e-02"/>
+      <inertia ixx="3.5232e-7" ixy="0.0" ixz="1.1744e-7" iyy="2.31944e-6" iyz="0" izz="2.23136e-6"/>
+    </inertial>
+  </link>
+  <link name="robotiq_85_right_knuckle_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/right_knuckle.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/right_knuckle.stl"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.01213197 0.0002 -0.00058647"/>
+      <mass value="1.384773208e-02"/>
+      <inertia ixx="3.5232e-7" ixy="0.0" ixz="-1.1744e-7" iyy="2.31944e-6" iyz="0.0" izz="2.23136e-6"/>
+    </inertial>
+  </link>
+  <link name="robotiq_85_left_finger_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/left_finger.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/left_finger.stl"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.00346899 -0.00079447 0.01867121"/>
+      <mass value="4.260376752e-02"/>
+      <inertia ixx="1.385792000000000e-05" ixy="0.0" ixz="-2.17264e-06" iyy="1.183208e-05" iyz="0.0" izz="5.19672e-06"/>
+    </inertial>
+  </link>
+  <link name="robotiq_85_right_finger_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/right_finger.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/right_finger.stl"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.00346899 -5.53e-06 0.01867121"/>
+      <mass value="4.260376752000000e-02"/>
+      <inertia ixx="1.385792e-05" ixy="0.0" ixz="2.17264e-06" iyy="1.183208e-05" iyz="0.0" izz="5.19672e-06"/>
+    </inertial>
+  </link>
+  <link name="robotiq_85_left_inner_knuckle_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/left_inner_knuckle.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/left_inner_knuckle.stl"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.01897699 0.00015001 0.02247101"/>
+      <mass value="2.969376448e-02"/>
+      <inertia ixx="9.57136e-06" ixy="0.0" ixz="-3.93424e-06" iyy="8.69056e-06" iyz="0.0" izz="8.19144e-06"/>
+    </inertial>
+  </link>
+  <link name="robotiq_85_right_inner_knuckle_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/right_inner_knuckle.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/right_inner_knuckle.stl"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.01926824 5.001e-05 0.02222178"/>
+      <mass value="2.969376448e-02"/>
+      <inertia ixx="9.42456e-06" ixy="0.0" ixz="3.9636e-06" iyy="8.69056e-06" iyz="0.0" izz="8.33824e-06"/>
+    </inertial>
+  </link>
+  <link name="robotiq_85_left_finger_tip_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/left_finger_tip.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/left_finger_tip.stl"/>
+      </geometry>
+      <surface>
+        <friction>
+          <ode>
+            <mu1>100000.0</mu1>
+            <mu2>100000.0</mu2>
+          </ode>
+        </friction>
+        <contact>
+          <ode>
+            <kp>1e+5</kp>
+            <kd>1</kd>
+            <soft_cfm>0</soft_cfm>
+            <soft_erp>0.2</soft_erp>
+            <minDepth>0.002</minDepth>
+            <maxVel>0</maxVel>
+          </ode>
+        </contact>
+      </surface>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="-0.01456706 -0.0008 0.01649701"/>
+      <mass value="4.268588744e-02"/>
+      <inertia ixx="1.048152e-05" ixy="0.0" ixz="3.5232e-6" iyy="1.197888e-05" iyz="0.0" izz="4.22784e-06"/>
+    </inertial>
+  </link>
+  <link name="robotiq_85_right_finger_tip_link">
+    <visual>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/visual/2f_85/right_finger_tip.dae"/>
+      </geometry>
+    </visual>
+    <collision>
+      <geometry>
+        <mesh filename="robotiq_2f85_meshes/collision/2f_85/right_finger_tip.stl"/>
+      </geometry>
+      <surface>
+        <friction>
+          <ode>
+            <mu1>100000.0</mu1>
+            <mu2>100000.0</mu2>
+          </ode>
+        </friction>
+        <contact>
+          <ode>
+            <kp>1e+5</kp>
+            <kd>1</kd>
+            <soft_cfm>0</soft_cfm>
+            <soft_erp>0.2</soft_erp>
+            <minDepth>0.002</minDepth>
+            <maxVel>0</maxVel>
+          </ode>
+        </contact>
+      </surface>
+    </collision>
+    <inertial>
+      <origin rpy="0 0 0" xyz="0.01456706 5e-05 0.01649701"/>
+      <mass value="4.268588744e-02"/>
+      <inertia ixx="1.048152e-05" ixy="0.0" ixz="-3.5232e-06" iyy="1.197888e-05" iyz="0.0" izz="4.22784e-06"/>
+    </inertial>
+  </link>
+  <joint name="robotiq_85_left_knuckle_joint" type="revolute">
+    <parent link="robotiq_85_base_link"/>
+    <child link="robotiq_85_left_knuckle_link"/>
+    <axis xyz="0 -1 0"/>
+    <origin rpy="0 0 0" xyz="0.03060114 0.0 0.05490452"/>
+    <limit effort="50" lower="0.0" upper="0.8" velocity="0.5"/>
+  </joint>
+  <joint name="robotiq_85_right_knuckle_joint" type="revolute">
+    <parent link="robotiq_85_base_link"/>
+    <child link="robotiq_85_right_knuckle_link"/>
+    <axis xyz="0 -1 0"/>
+    <origin rpy="0 0 0" xyz="-0.03060114 0.0 0.05490452"/>
+    <limit effort="50" lower="-0.8" upper="0.0" velocity="0.5"/>
+  </joint>
+  <joint name="robotiq_85_left_finger_joint" type="fixed">
+    <parent link="robotiq_85_left_knuckle_link"/>
+    <child link="robotiq_85_left_finger_link"/>
+    <origin rpy="0 0 0" xyz="0.03152616 0.0 -0.00376347"/>
+  </joint>
+  <joint name="robotiq_85_right_finger_joint" type="fixed">
+    <parent link="robotiq_85_right_knuckle_link"/>
+    <child link="robotiq_85_right_finger_link"/>
+    <origin rpy="0 0 0" xyz="-0.03152616 0.0 -0.00376347"/>
+  </joint>
+  <joint name="robotiq_85_left_inner_knuckle_joint" type="revolute">
+    <parent link="robotiq_85_base_link"/>
+    <child link="robotiq_85_left_inner_knuckle_link"/>
+    <axis xyz="0 -1 0"/>
+    <origin rpy="0 0 0" xyz="0.0127 0.0 0.06142"/>
+    <limit effort="50" lower="0" upper="0.8" velocity="0.5"/>
+    
+  </joint>
+  <joint name="robotiq_85_right_inner_knuckle_joint" type="revolute">
+    <parent link="robotiq_85_base_link"/>
+    <child link="robotiq_85_right_inner_knuckle_link"/>
+    <axis xyz="0 -1 0"/>
+    <origin rpy="0 0 0" xyz="-0.0127 0.0 0.06142"/>
+    <limit effort="50" lower="-0.8" upper="0" velocity="0.5"/>
+  </joint>
+  <joint name="robotiq_85_left_finger_tip_joint" type="revolute">
+    <parent link="robotiq_85_left_finger_link"/>
+    <child link="robotiq_85_left_finger_tip_link"/>
+    <axis xyz="0 -1 0"/>
+    <origin rpy="0 0 0" xyz="0.00563134 0.0 0.04718515"/>
+    <limit effort="50" lower="-0.8" upper="0" velocity="0.5"/>
+  </joint>
+  <joint name="robotiq_85_right_finger_tip_joint" type="revolute">
+    <parent link="robotiq_85_right_finger_link"/>
+    <child link="robotiq_85_right_finger_tip_link"/>
+    <axis xyz="0 -1 0"/>
+    <origin rpy="0 0 0" xyz="-0.00563134 0.0 0.04718515"/>
+    <limit effort="50" lower="0" upper="0.8" velocity="0.5"/>
+  </joint>
+  
+  <link name="gripper_center"/>
+
+  <joint name="ee_fixed_joint" type="fixed">
+    <parent link="robotiq_85_base_link"/>
+    <child link="gripper_center"/>
+    <origin rpy="0 0.0 -1.57079632679" xyz="0.0 0. 0.1635"/>
+  </joint>
+
+</robot>
+

replay_init_dependencies/gen_traj.py

@@ -0,0 +1,339 @@
+import numpy as np
+import open3d as o3d
+from scipy.spatial.transform import Rotation as R
+
+
+def get_rotation(right_vector, up_vector, look_at_vector, z_front=True):
+    rotation = np.eye(3)
+    rotation[:, 0] = right_vector
+    rotation[:, 1] = up_vector
+    rotation[:, 2] = look_at_vector
+    if not z_front:
+        rotation = rotation @ R.from_euler('y', 180, degrees=True).as_matrix()
+    return rotation
+
+
+class TrajDescription:
+    RRot: str = "r_rot"
+    LRot: str = "l_rot"
+    UpRot: str = "up_rot"
+    DownRot: str = "down_rot"
+    Front: str = "front"
+    Back: str = "back"
+    Left: str = "left"
+    Right: str = "right"
+    Up: str = "up"
+    Down: str = "down"
+
+    def __init__(self, start_pt, look_at_pt, description_list):
+        self.z_front = False
+        self.up_vector = np.array([0., 0., -1.])
+        self.init_pose = self.get_pose(start_pt, look_at_pt, self.up_vector)
+        self.description_list = description_list
+        self.poses = self.generate_poses()
+
+    def get_pose(self, start_pt, look_at_pt, up_vector):
+        look_at_vector = look_at_pt - start_pt
+        look_at_vector = look_at_vector / np.linalg.norm(look_at_vector)
+        up_vector = up_vector
+        right_vector = np.cross(look_at_vector, up_vector)
+        right_vector = right_vector / np.linalg.norm(right_vector)
+        up_vector = np.cross(look_at_vector, right_vector)
+        up_vector = up_vector / np.linalg.norm(up_vector)
+
+        pose = np.eye(4)
+        pose[:3, :3] = get_rotation(right_vector, up_vector, look_at_vector, self.z_front)
+        pose[:3, 3] = start_pt
+        return pose
+
+    def generate_poses(self):
+        poses = [self.init_pose]
+        current_pose = self.init_pose.copy()
+
+        for description in self.description_list:
+            pose_type = description["type"]
+            value = description["value"]
+
+            if pose_type == self.RRot:
+                rotation_matrix = R.from_euler('y', -value, degrees=True).as_matrix()
+                new_pose = current_pose.copy()
+                new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
+                poses.append(new_pose)
+                poses.append(current_pose)
+
+            elif pose_type == self.LRot:
+                rotation_matrix = R.from_euler('y', value, degrees=True).as_matrix()
+                new_pose = current_pose.copy()
+                new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
+                poses.append(new_pose)
+                poses.append(current_pose)
+
+            elif pose_type == self.UpRot:
+                rotation_matrix = R.from_euler('x', -value, degrees=True).as_matrix()
+                new_pose = current_pose.copy()
+                new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
+                poses.append(new_pose)
+                poses.append(current_pose)
+
+            elif pose_type == self.DownRot:
+                rotation_matrix = R.from_euler('x', value, degrees=True).as_matrix()
+                new_pose = current_pose.copy()
+                new_pose[:3, :3] = current_pose[:3, :3] @ rotation_matrix
+                poses.append(new_pose)
+                poses.append(current_pose)
+
+            elif pose_type == self.Front:
+                new_pose = current_pose.copy()
+                new_pose[1, 3] = current_pose[1, 3] + value
+                poses.append(new_pose)
+                current_pose = new_pose
+
+            elif pose_type == self.Back:
+                new_pose = current_pose.copy()
+                new_pose[1, 3] = current_pose[1, 3] - value
+                poses.append(new_pose)
+                current_pose = new_pose
+
+            elif pose_type == self.Left:
+                new_pose = current_pose.copy()
+                new_pose[0, 3] = current_pose[0, 3] - value
+                poses.append(new_pose)
+                current_pose = new_pose
+
+            elif pose_type == self.Right:
+                new_pose = current_pose.copy()
+                new_pose[0, 3] = current_pose[0, 3] + value
+                poses.append(new_pose)
+                current_pose = new_pose
+
+            elif pose_type == self.Up:
+                new_pose = current_pose.copy()
+                new_pose[2, 3] = current_pose[2, 3] + value
+                poses.append(new_pose)
+                current_pose = new_pose
+
+            elif pose_type == self.Down:
+                new_pose = current_pose.copy()
+                new_pose[2, 3] = current_pose[2, 3] - value
+                poses.append(new_pose)
+                current_pose = new_pose
+
+            else:
+                raise ValueError(f"Invalid description: {description}")
+
+        return poses
+
+
+class TrajGenerator:
+    def __init__(self, view_num: int, look_at_pts: np.ndarray, up_vectors: np.ndarray, start_pts: np.ndarray,
+                 reverse_order=False, z_front=True):
+        """Generates a circular trajectory around a specified point.
+        Args:
+            view_num (int): Number of views/poses to generate.
+            look_at_pts (np.ndarray): The point(s) to look at, shape (N, 3).
+            up_vectors (np.ndarray): The up direction(s), shape (N, 3).
+            start_pts (np.ndarray): The starting point(s) of the trajectory, shape (N, 3).
+            reverse_order (bool): If True, reverse the order of generated poses.
+            z_front (bool): If True, the camera's Z-axis points forward; otherwise, it points backward.
+        """
+        assert look_at_pts.shape[0] == up_vectors.shape[0] == start_pts.shape[
+            0], "Input arrays must have the same number of points."
+        self.view_num = view_num
+        self.look_at_pts = look_at_pts
+        self.up_vectors = up_vectors
+        self.start_pts = start_pts
+        self.z_front = z_front
+        self.poses = self.generate_poses(reverse_order)
+        self.interpolated_poses = self.interpolate_poses(self.poses, num_interpolations=self.view_num, method='slerp')
+
+    def generate_poses(self, reverse_order=False):
+        poses = []
+        length = self.look_at_pts.shape[0]
+        for i in range(length):
+            order = i if not reverse_order else length - i - 1
+            pose = self.get_pose(self.start_pts[order], self.look_at_pts[order], self.up_vectors[order])
+            poses.append(pose)
+        return poses
+
+    def get_pose(self, start_pt, look_at_pt, up_vector, with_noise=False):
+        look_at_vector = look_at_pt - start_pt
+        look_at_vector = look_at_vector / np.linalg.norm(look_at_vector)
+        up_vector = up_vector
+        if with_noise:
+            noise = np.random.uniform(0, 0.1, up_vector.shape)
+            right_vector = np.cross(look_at_vector, up_vector + noise)
+        else:
+            right_vector = np.cross(look_at_vector, up_vector)
+        right_vector = right_vector / np.linalg.norm(right_vector)
+        up_vector = np.cross(look_at_vector, right_vector)
+        up_vector = up_vector / np.linalg.norm(up_vector)
+
+        pose = np.eye(4)
+        pose[:3, :3] = self.get_rotation(right_vector, up_vector, look_at_vector, self.z_front)
+        pose[:3, 3] = start_pt
+        return pose
+
+    def get_rotation(self, right_vector, up_vector, look_at_vector, z_front=True):
+        rotation = np.eye(3)
+        rotation[:, 0] = right_vector
+        rotation[:, 1] = up_vector
+        rotation[:, 2] = look_at_vector
+        if not z_front:
+            rotation = rotation @ R.from_euler('y', 180, degrees=True).as_matrix()
+        return rotation
+
+    def interpolate_poses(self, poses, num_interpolations, method='slerp'):
+        assert method in ['slerp', 'linear'], "Method must be 'slerp' or 'linear'."
+        return poses
+
+
+class CircularTrajGenerator(TrajGenerator):
+    def __init__(self, view_num, look_at_pt, top_down_dist, top_down_direrction, radius, reverse_order=False,
+                 z_front=True):
+        start_pt = look_at_pt + np.array([radius, 0, top_down_dist])
+        look_at_pts = np.tile(look_at_pt, (view_num, 1))
+        up_vectors = np.tile(top_down_direrction, (view_num, 1))
+        start_pts = np.zeros((view_num, 3))
+        for i in range(view_num):
+            angle = np.pi * 2 * i / view_num
+            start_point_x = look_at_pt[0] + radius * np.cos(angle)
+            start_point_y = look_at_pt[1] + radius * np.sin(angle)
+            start_point_z = look_at_pt[2] + top_down_dist
+            start_pts[i] = np.array([start_point_x, start_point_y, start_point_z])
+        super().__init__(view_num, look_at_pts, up_vectors, start_pts, reverse_order, z_front)
+
+
+class TrajVisualizer:
+    def __init__(self, poses: list) -> None:
+        """Visualizes SE3 poses."""
+        self.poses = poses
+
+    def visualize(self):
+        geometries = []
+        for pose in self.poses:
+            mesh_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.1, origin=[0, 0, 0])
+            mesh_frame.transform(pose)
+            geometries.append(mesh_frame)
+        o3d.visualization.draw_geometries(geometries)
+
+
+def generate_camera_data_from_poses(poses, name="unknown"):
+    camera_data = {
+        "name": name,
+        "width": 640,
+        "height": 480,
+        "focus_length": 13,
+        "horizontal_aperture": 15.2908,
+        "vertical_aperture": 15.2908,
+        "is_local": False,
+        "trajectory": []
+    }
+
+    for pose in poses:
+        position = pose[:3, 3].tolist()
+        rotation_matrix = pose[:3, :3]
+        quaternion_xyzw = R.from_matrix(rotation_matrix).as_quat().tolist()
+        quaternion_wxyz = [quaternion_xyzw[3], quaternion_xyzw[0], quaternion_xyzw[1], quaternion_xyzw[2]]
+
+        camera_data["trajectory"].append({
+            "position": position,
+            "quaternion": quaternion_wxyz
+        })
+
+    return camera_data
+
+
+def generate_circular_trajectory(top_down_dist=0.6, radius=0.3, look_at_pt=(0, 0, 0)):
+    VIEW_NUM = 100
+    LOOK_AT_PT = np.array(look_at_pt)
+    TOP_DOWN_DIST = [top_down_dist]
+    TOP_DOWN_DIRECTION = np.array([0., 0., -1.])
+    RADIUS = [radius]
+
+    pose_generator = CircularTrajGenerator(VIEW_NUM, LOOK_AT_PT, TOP_DOWN_DIST[0], TOP_DOWN_DIRECTION, RADIUS[0],
+                                           reverse_order=False, z_front=False)
+    return pose_generator.poses
+
+
+def generate_trajectory_from_description(start_pt, look_at_pt, description_list):
+    traj_description = TrajDescription(start_pt, look_at_pt, description_list)
+    return traj_description.poses
+
+
+def generate_front_left_trajectory():
+    VIEW_NUM = 100
+    LOOK_AT_PT = np.array([-0.717877984046936, -2.1181700229644775, 0.7807716727256775])
+    TOP_DOWN_DIST = [0.6]
+    TOP_DOWN_DIRECTION = np.array([0., 0., -1.])
+    RADIUS = [0.3]
+
+    start_pt = LOOK_AT_PT + np.array([RADIUS[0], 0, TOP_DOWN_DIST[0]])
+    look_at_pts = np.tile(LOOK_AT_PT, (VIEW_NUM, 1))
+    up_vectors = np.tile(TOP_DOWN_DIRECTION, (VIEW_NUM, 1))
+    start_pts = np.zeros((VIEW_NUM, 3))
+    for i in range(VIEW_NUM):
+        angle = np.pi * 2 * i / VIEW_NUM + np.pi / 4
+        start_point_x = LOOK_AT_PT[0] + RADIUS[0] * np.cos(angle)
+        start_point_y = LOOK_AT_PT[1] + RADIUS[0] * np.sin(angle)
+        start_point_z = LOOK_AT_PT[2] + TOP_DOWN_DIST[0]
+        start_pts[i] = np.array([start_point_x, start_point_y, start_point_z])
+
+    pose_generator = TrajGenerator(VIEW_NUM, look_at_pts, up_vectors, start_pts, reverse_order=False, z_front=False)
+    return pose_generator.poses
+
+
+def main():
+    workspace_size = np.array([0.816116, 0.675734, 0.3])
+    workspace_position = np.array([-0.717877984046936, -2.1181700229644775, 0.7807716727256775])
+    desc_1 = [
+        # {"type": TrajDescription.LRot, "value": 20},
+        # {"type": TrajDescription.RRot, "value": 20},
+        # {"type": TrajDescription.DownRot, "value": 20},
+        # {"type": TrajDescription.UpRot, "value": 20},
+        {"type": TrajDescription.Left, "value": workspace_size[0]},
+        {"type": TrajDescription.Front, "value": workspace_size[1]},
+        # {"type": TrajDescription.Back, "value": workspace_size[1]},
+        # {"type": TrajDescription.Right, "value": workspace_size[0]}
+    ]
+    desc_2 = [
+        {"type": TrajDescription.LRot, "value": 20},
+        {"type": TrajDescription.RRot, "value": 20},
+        {"type": TrajDescription.DownRot, "value": 20},
+        {"type": TrajDescription.UpRot, "value": 20},
+        {"type": TrajDescription.Right, "value": workspace_size[0]},
+        {"type": TrajDescription.Front, "value": workspace_size[1]},
+        {"type": TrajDescription.Back, "value": workspace_size[1]},
+        {"type": TrajDescription.Left, "value": workspace_size[0]}
+    ]
+    desc_3 = [
+        {"type": TrajDescription.DownRot, "value": 20},
+        {"type": TrajDescription.UpRot, "value": 20},
+        {"type": TrajDescription.Left, "value": workspace_size[0]},
+        {"type": TrajDescription.Front, "value": workspace_size[1]},
+        {"type": TrajDescription.RRot, "value": 40},
+        {"type": TrajDescription.LRot, "value": 40},
+        {"type": TrajDescription.Back, "value": workspace_size[1]},
+        {"type": TrajDescription.Right, "value": workspace_size[0]}
+    ]
+    desc_4 = [
+        {"type": TrajDescription.DownRot, "value": 20},
+        {"type": TrajDescription.UpRot, "value": 20},
+        {"type": TrajDescription.Right, "value": workspace_size[0]},
+        {"type": TrajDescription.Front, "value": workspace_size[1]},
+        {"type": TrajDescription.LRot, "value": 40},
+        {"type": TrajDescription.RRot, "value": 40},
+        {"type": TrajDescription.Back, "value": workspace_size[1]},
+        {"type": TrajDescription.Left, "value": workspace_size[0]}
+    ]
+
+    start_pt = workspace_position - np.array([0, workspace_size[1] / 2, -0.3])
+    poses_1 = generate_trajectory_from_description(start_pt, workspace_position, desc_1)
+    poses_2 = generate_trajectory_from_description(start_pt, workspace_position, desc_2)
+    poses_3 = generate_trajectory_from_description(start_pt, workspace_position, desc_3)
+    poses_4 = generate_trajectory_from_description(start_pt, workspace_position, desc_4)
+    visualizer = TrajVisualizer(poses_1)
+    visualizer.visualize()
+
+
+if __name__ == "__main__":
+    main()

runners/data_generator.py

@@ -1,11 +1,23 @@
+import json
+
 from pyboot import stereotype
 from pyboot.runner import Runner
-
+from pyboot.utils.log import Log
 from data_gen_dependencies.data_generate import generate_data
+
 @stereotype.runner("data_generator")
 class DataGenerator(Runner):
     def __init__(self, config_path: str):
         super().__init__(config_path)
+        self.generate_config = self.config["generate"]
+        self.input_target_task_path = self.generate_config["input_target_task_path"]
+        self.input_data_root = self.generate_config["input_data_root"]
+        self.server_url = self.generate_config["server_url"]
+        self.target_tasks = json.load(open(self.input_target_task_path, "r"))
 
     def run(self):
-        pass
+        for task_template_name, task_list in self.target_tasks.items():
+            Log.info(f"Generating from template: {task_template_name} | tasks number: {len(task_list)}")
+            generate_data(task_list, self.server_url)
+            Log.warning("multi task templates generation is not implemented yet, stop after one template generation.")
+            break

runners/data_replayer.py

@@ -1,8 +1,8 @@
 from pyboot import stereotype
 from pyboot.runner import Runner
 
-@stereotype.runner("data_recorder")
-class DataRecorder(Runner):
+@stereotype.runner("data_replayer")
+class DataReplayer(Runner):
     def __init__(self, config_path: str):
         super().__init__(config_path)
 

runners/replay_initializer.py

@@ -0,0 +1,187 @@
+import os
+import json
+import numpy as np
+from dataclasses import dataclass
+from pyboot.utils.log import Log
+from replay_init_dependencies.gen_traj import generate_circular_trajectory, generate_camera_data_from_poses, \
+    generate_trajectory_from_description, TrajDescription
+
+
+@dataclass
+class ReplayConfig:
+    replay_dir: str = ""
+    recording_data_dir: str = ""
+    task_template_file: str = ""
+    use_origin_camera: bool = False
+    use_trajectory_camera: bool = True
+    space_interpolation: bool = False
+
+
+class ReplayInitializer:
+    TRAJECTORY_FILE_NAME: str = "camera_trajectory.json"
+
+    def __init__(self, replay_config: ReplayConfig):
+        self.replay_config = replay_config
+
+    def init_replay(self):
+        if not os.path.exists(self.replay_config.replay_dir):
+            os.makedirs(self.replay_config.replay_dir)
+        recording_data_dict = self.load_recording_data()
+        data_length = len(recording_data_dict)
+        task_template = self.load_task_template(data_length)
+        self.workspace = self.load_workspace(task_template)
+        self.generate_trajectory_file()
+        self.generate_replay_input(task_template, recording_data_dict)
+        task_template_file = os.path.join(self.replay_config.replay_dir,
+                                          os.path.basename(self.replay_config.task_template_file))
+        command = f"omni_python main.py --task_template {task_template_file} --use_recording"
+
+        Log.success(f"Replay initialized successfully. You can now run the replay by running following command:")
+        Log.success(f"{command}")
+
+    def load_workspace(self, task_template):
+        workspace_name = task_template["scene"]["scene_id"].split("/")[-1]
+        workspace_info = task_template["scene"]["function_space_objects"][workspace_name]
+        return workspace_info
+
+    def load_recording_data(self):
+        recording_data_dict = {}
+        for data_dir in os.listdir(self.replay_config.recording_data_dir):
+            if os.path.isdir(os.path.join(self.replay_config.recording_data_dir, data_dir)):
+                Log.info(f"Loading recording data from {data_dir}")
+                state = None
+                task_result = None
+                for file in os.listdir(os.path.join(self.replay_config.recording_data_dir, data_dir)):
+                    if file == "state.json":
+                        state = json.load(open(os.path.join(self.replay_config.recording_data_dir, data_dir, file)))
+                    elif file == "task_result.json":
+                        task_result = json.load(
+                            open(os.path.join(self.replay_config.recording_data_dir, data_dir, file)))
+                if state is None or task_result is None:
+                    Log.warning(f"state or task result not found in {data_dir}")
+                else:
+                    recording_data_dict[data_dir] = {
+                        "idx": int(data_dir.split('_')[-1][:-1]) + 1,
+                        "state": state,
+                        "task_result": task_result
+                    }
+        Log.success(f"Recording data loaded: {len(recording_data_dict)} data in total.")
+        return recording_data_dict
+
+    def load_task_template(self, data_length):
+        task_template = json.load(open(self.replay_config.task_template_file))
+        replay_config = {
+            "use_origin_camera": self.replay_config.use_origin_camera,
+            "use_trajectory_camera": self.replay_config.use_trajectory_camera,
+            "replay_dir": self.replay_config.replay_dir,
+            "episodes": data_length,
+            "cam_trajectory_file": os.path.join(self.replay_config.replay_dir, self.TRAJECTORY_FILE_NAME),
+            "space_interpolation": self.replay_config.space_interpolation
+        }
+        task_template["replay"] = replay_config
+        Log.success(f"Task template loaded successfully. replay config: {replay_config}")
+        return task_template
+
+    def generate_trajectory_file(self):
+        workspace_position = self.workspace["position"]
+        workspace_size = self.workspace["size"]
+        trajectory_data = []
+        # poses_mid = generate_circular_trajectory(0.5,0.3,workspace_position)
+        # poses_low = generate_circular_trajectory(0.05,0.5,workspace_position)
+        # poses_high = generate_circular_trajectory(0.7,0.15,workspace_position)
+        desc_1 = [
+            {"type": TrajDescription.LRot, "value": 20},
+            {"type": TrajDescription.RRot, "value": 20},
+            {"type": TrajDescription.DownRot, "value": 20},
+            {"type": TrajDescription.UpRot, "value": 20},
+            {"type": TrajDescription.Left, "value": workspace_size[0] / 2},
+            {"type": TrajDescription.Front, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.Back, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.Right, "value": workspace_size[0] / 2}
+        ]
+        desc_2 = [
+            {"type": TrajDescription.LRot, "value": 20},
+            {"type": TrajDescription.RRot, "value": 20},
+            {"type": TrajDescription.DownRot, "value": 20},
+            {"type": TrajDescription.UpRot, "value": 20},
+            {"type": TrajDescription.Right, "value": workspace_size[0] / 2},
+            {"type": TrajDescription.Front, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.Back, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.Left, "value": workspace_size[0] / 2}
+        ]
+        desc_3 = [
+            {"type": TrajDescription.DownRot, "value": 20},
+            {"type": TrajDescription.UpRot, "value": 20},
+            {"type": TrajDescription.Left, "value": workspace_size[0] / 2},
+            {"type": TrajDescription.Front, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.RRot, "value": 75},
+            {"type": TrajDescription.LRot, "value": 75},
+            {"type": TrajDescription.Back, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.Right, "value": workspace_size[0] / 2}
+        ]
+        desc_4 = [
+            {"type": TrajDescription.DownRot, "value": 20},
+            {"type": TrajDescription.UpRot, "value": 20},
+            {"type": TrajDescription.Right, "value": workspace_size[0] / 2},
+            {"type": TrajDescription.Front, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.LRot, "value": 75},
+            {"type": TrajDescription.RRot, "value": 75},
+            {"type": TrajDescription.Back, "value": workspace_size[1] / 2},
+            {"type": TrajDescription.Left, "value": workspace_size[0] / 2}
+        ]
+
+        start_pt = workspace_position - np.array([0, workspace_size[1], -0.3])
+        workspace_bottom_center = workspace_position - np.array([0, 0, workspace_size[2] / 2])
+        poses_1 = generate_trajectory_from_description(start_pt, workspace_bottom_center, desc_1)
+        poses_2 = generate_trajectory_from_description(start_pt, workspace_bottom_center, desc_2)
+        poses_3 = generate_trajectory_from_description(start_pt, workspace_bottom_center, desc_3)
+        poses_4 = generate_trajectory_from_description(start_pt, workspace_bottom_center, desc_4)
+
+        camera_data_1 = generate_camera_data_from_poses(poses_1, "cam_1")
+        camera_data_2 = generate_camera_data_from_poses(poses_2, "cam_2")
+        camera_data_3 = generate_camera_data_from_poses(poses_3, "cam_3")
+        camera_data_4 = generate_camera_data_from_poses(poses_4, "cam_4")
+
+        trajectory_data = [camera_data_1, camera_data_2, camera_data_3, camera_data_4]
+        trajectory_path = os.path.join(self.replay_config.replay_dir, self.TRAJECTORY_FILE_NAME)
+        with open(trajectory_path, "w") as f:
+            json.dump(trajectory_data, f)
+        Log.success(f"Trajectory file generated successfully. Trajectory file saved to {trajectory_path}")
+
+    def generate_trajectory_file_from_json(self):
+        workspace_position = self.workspace["position"]
+        trajectory_data = []
+        poses_mid = generate_circular_trajectory(0.5, 0.3, workspace_position)
+        poses_low = generate_circular_trajectory(0.05, 0.5, workspace_position)
+        poses_high = generate_circular_trajectory(0.7, 0.15, workspace_position)
+        camera_data_mid = generate_camera_data_from_poses(poses_mid, "cam_mid")
+        camera_data_low = generate_camera_data_from_poses(poses_low, "cam_low")
+        camera_data_high = generate_camera_data_from_poses(poses_high, "cam_high")
+        trajectory_data = [camera_data_mid, camera_data_low, camera_data_high]
+        trajectory_path = os.path.join(self.replay_config.replay_dir, self.TRAJECTORY_FILE_NAME)
+        with open(trajectory_path, "w") as f:
+            json.dump(trajectory_data, f)
+        Log.success(f"Trajectory file generated successfully. Trajectory file saved to {trajectory_path}")
+
+    def generate_replay_input(self, task_template, recording_data_dict):
+        states_dir = os.path.join(self.replay_config.replay_dir, "states")
+        results_dir = os.path.join(self.replay_config.replay_dir, "results")
+        if not os.path.exists(states_dir):
+            os.makedirs(states_dir)
+        if not os.path.exists(results_dir):
+            os.makedirs(results_dir)
+        for data_dir in recording_data_dict:
+            state = recording_data_dict[data_dir]["state"]
+            task_result = recording_data_dict[data_dir]["task_result"]
+            state_file = os.path.join(states_dir, f"{recording_data_dict[data_dir]['idx']}.state.json")
+            result_file = os.path.join(results_dir, f"{recording_data_dict[data_dir]['idx']}.task_result.json")
+            with open(state_file, "w") as f:
+                json.dump(state, f)
+            with open(result_file, "w") as f:
+                json.dump(task_result, f)
+        with open(os.path.join(self.replay_config.replay_dir, os.path.basename(self.replay_config.task_template_file)),
+                  "w") as f:
+            json.dump(task_template, f)
+        Log.success(f"Replay input generated: {len(recording_data_dict)} data in total.")
+        Log.success(
+            f"Replay input saved to '{self.replay_config.replay_dir}', new task file is saved as '{os.path.basename(self.replay_config.task_template_file)}' in '{self.replay_config.replay_dir}'")

runners/task_generator.py

@@ -29,8 +29,8 @@ class TaskGenerator(Runner):
             task_template = json.load(open(task_template_path, "r"))
             generated_tasks_path_list = self.generate_from_template(task_template_name, task_template)
             Log.success(f"Generated {len(generated_tasks_path_list)} tasks from <{task_template_name}>")
-            generate_results["task_template_name"] = generated_tasks_path_list
-        json.dump(generate_results, open(self.output_generate_result_path, "w"))
+            generate_results[task_template_name] = generated_tasks_path_list
+            json.dump(generate_results, open(self.output_generate_result_path, "w"))
 
 
     def generate_from_template(self, template_name: str, template: dict) -> List[str]:

task_gen_dependencies/object.py

@@ -48,9 +48,9 @@ class OmniObject:
                         "direction": np.array([-1,0,0.3]),
                     },
                 ],
-                'click': {
+                'press': {
                     "part": "finger edge",
-                    "task": "click",
+                    "task": "press",
                     "xyz": np.array([0,0,0]),
                     "direction": np.array([0,0,1]),
                 },
@@ -181,16 +181,14 @@ class OmniObject:
                             grasp_data['grasp_pose'][:, :3, 3] = translation
                             grasp_data['width'] = width
 
-                            print(f"Grasp num of {obj.name} after NMS: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
+                            #print(f"Grasp num of {obj.name} after NMS: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
                             # downsample
                             if grasp_data['grasp_pose'].shape[0]>max_grasps:
                                 grasp_num = grasp_data['grasp_pose'].shape[0]
                                 index = np.random.choice(grasp_num, max_grasps, replace=False)
                                 grasp_data['grasp_pose'] = grasp_data['grasp_pose'][index]
                                 grasp_data['width'] = grasp_data['width'][index]
-                                print(f"Grasp num of {obj.name} after downsample: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
-                        else:
-                            print(f"Grasp num of {obj.name}: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
+                                #print(f"Grasp num of {obj.name} after downsample: {grasp_data['grasp_pose'].shape[0]}/{N_grasp}")
 
                         action_info[grasp_part] = grasp_data
 

task_gen_dependencies/task_generate.py

@@ -99,7 +99,7 @@ class OldTaskGenerator:
     def __init__(self, task_template, data_root):
         self.data_root = data_root
         self.init_info(task_template)
-        self.task_template = task_template
+
 
     def _load_json(self, relative_path):
         with open(os.path.join(self.data_root, relative_path), 'r') as file:
@@ -171,15 +171,15 @@ class OldTaskGenerator:
         scene_info = task_template["scene"]
         self.scene_usd = task_template["scene"]["scene_usd"]
         self.task_template = {
-            "scene_usd": self.scene_usd,
-            "arm": arm,
+            "scene": task_template["scene"],
             "task_name": task_template["task"],
-            "robot_id": robot_id,
+            "robot":{"robot_id": robot_id,"arm": arm},
             "stages": task_template['stages'],
             "object_with_material": task_template.get('object_with_material', {}),
             "lights": task_template.get('lights', {}),
             "cameras": task_template.get('cameras', {}),
-            "objects": []
+            "objects": [],
+            "recording_setting": task_template.get('recording_setting', {})
         }
         constraint = task_template.get("constraints")
         robot_init_workspace_id = scene_info["scene_id"].split('/')[-1]
@@ -232,14 +232,10 @@ class OldTaskGenerator:
             generated_tasks_path.append(output_file)
             task_instance['objects'] = []
             task_instance['objects'] += self.fix_obj_infos
-            task_instance['robot'] = {
-                "init_position" : self.robot_init_pose["position"],
-                "init_rotation" : self.robot_init_pose["quaternion"],
-                "robot_cfg": self.robot_cfg,
-                "stand": self.stand,
-            }
-            task_instance['recording_setting'] = self.task_template["recording_setting"]
-
+            task_instance['robot']["init_position"] = self.robot_init_pose["position"]
+            task_instance['robot']["init_rotation"] = self.robot_init_pose["quaternion"]
+            task_instance['robot']["robot_cfg"] = self.robot_cfg
+            task_instance['robot']["stand"] = self.stand
             flag_failed = False
             for key in self.layouts:
                 obj_infos = self.layouts[key]()