release repository

examples/isaac_sim/helper.py

@@ -0,0 +1,134 @@
+#
+# Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+#
+
+# Standard Library
+from typing import Dict, List
+
+# Third Party
+import numpy as np
+from matplotlib import cm
+from omni.isaac.core import World
+from omni.isaac.core.materials import OmniPBR
+from omni.isaac.core.objects import cuboid
+from omni.isaac.core.robots import Robot
+
+try:
+    # Third Party
+    from omni.isaac.urdf import _urdf  # isaacsim 2022.2
+except ImportError:
+    from omni.importer.urdf import _urdf  # isaac sim 2023.1
+# CuRobo
+from curobo.util_file import get_assets_path, get_filename, get_path_of_dir, join_path
+
+
+############################################################
+def add_robot_to_scene(
+    robot_config: Dict,
+    my_world: World,
+    load_from_usd: bool = False,
+    subroot: str = "",
+    robot_name: str = "robot",
+    position: np.array = np.array([0, 0, 0]),
+):
+    urdf_interface = _urdf.acquire_urdf_interface()
+
+    import_config = _urdf.ImportConfig()
+    import_config.merge_fixed_joints = False
+    import_config.convex_decomp = False
+    import_config.import_inertia_tensor = True
+    import_config.fix_base = True
+    import_config.make_default_prim = False
+    import_config.self_collision = False
+    import_config.create_physics_scene = True
+    import_config.import_inertia_tensor = False
+    import_config.default_drive_strength = 20000
+    import_config.default_position_drive_damping = 500
+    import_config.default_drive_type = _urdf.UrdfJointTargetType.JOINT_DRIVE_POSITION
+    import_config.distance_scale = 1
+    import_config.density = 0.0
+    full_path = join_path(get_assets_path(), robot_config["kinematics"]["urdf_path"])
+    robot_path = get_path_of_dir(full_path)
+    filename = get_filename(full_path)
+    imported_robot = urdf_interface.parse_urdf(robot_path, filename, import_config)
+    dest_path = subroot
+    robot_path = urdf_interface.import_robot(
+        robot_path,
+        filename,
+        imported_robot,
+        import_config,
+        dest_path,
+    )
+
+    # prim_path = omni.usd.get_stage_next_free_path(
+    # my_world.scene.stage, str(my_world.scene.stage.GetDefaultPrim().GetPath()) + robot_path, False)
+    # print(prim_path)
+    # robot_prim = my_world.scene.stage.OverridePrim(prim_path)
+    # robot_prim.GetReferences().AddReference(dest_path)
+
+    robot = my_world.scene.add(
+        Robot(
+            prim_path=robot_path,
+            name=robot_name,
+            position=position,
+        )
+    )
+
+    return robot, robot_path
+
+
+class VoxelManager:
+    def __init__(
+        self,
+        num_voxels: int = 5000,
+        size: float = 0.02,
+        color: List[float] = [1, 1, 1],
+        prefix_path: str = "/World/curobo/voxel_",
+        material_path: str = "/World/looks/v_",
+    ) -> None:
+        self.cuboid_list = []
+        self.cuboid_material_list = []
+        self.disable_idx = num_voxels
+        for i in range(num_voxels):
+            target_material = OmniPBR("/World/looks/v_" + str(i), color=np.ravel(color))
+
+            cube = cuboid.VisualCuboid(
+                prefix_path + str(i),
+                position=np.array([0, 0, -10]),
+                orientation=np.array([1, 0, 0, 0]),
+                size=size,
+                visual_material=target_material,
+            )
+            self.cuboid_list.append(cube)
+            self.cuboid_material_list.append(target_material)
+            cube.set_visibility(True)
+
+    def update_voxels(self, voxel_position: np.ndarray, color_axis: int = 0):
+        max_index = min(voxel_position.shape[0], len(self.cuboid_list))
+
+        jet = cm.get_cmap("hot")  # .reversed()
+        z_val = voxel_position[:, 0]
+
+        jet_colors = jet(z_val)
+
+        for i in range(max_index):
+            self.cuboid_list[i].set_visibility(True)
+
+            self.cuboid_list[i].set_local_pose(translation=voxel_position[i])
+            self.cuboid_material_list[i].set_color(jet_colors[i][:3])
+
+        for i in range(max_index, len(self.cuboid_list)):
+            self.cuboid_list[i].set_local_pose(translation=np.ravel([0, 0, -10.0]))
+
+            # self.cuboid_list[i].set_visibility(False)
+
+    def clear(self):
+        for i in range(len(self.cuboid_list)):
+            self.cuboid_list[i].set_local_pose(translation=np.ravel([0, 0, -10.0]))