finish data generate

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,