release repository

examples/isaac_sim/load_all_robots.py

@@ -0,0 +1,180 @@
+#
+# 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.
+#
+
+# Third Party
+import torch
+
+a = torch.zeros(4, device="cuda:0")
+
+# Standard Library
+
+# Standard Library
+import argparse
+
+# Third Party
+from omni.isaac.kit import SimulationApp
+
+# CuRobo
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
+
+# from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
+from curobo.geom.sdf.world import CollisionCheckerType
+from curobo.geom.types import WorldConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.math import Pose
+from curobo.types.robot import RobotConfig
+from curobo.types.state import JointState
+from curobo.util_file import (
+    get_motion_gen_robot_list,
+    get_robot_configs_path,
+    get_robot_path,
+    get_world_configs_path,
+    join_path,
+    load_yaml,
+)
+from curobo.wrap.model.robot_world import RobotWorld, RobotWorldConfig
+from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--headless", action="store_true", help="When True, enables headless mode", default=False
+)
+parser.add_argument(
+    "--visualize_spheres",
+    action="store_true",
+    help="When True, visualizes robot spheres",
+    default=False,
+)
+
+args = parser.parse_args()
+simulation_app = SimulationApp(
+    {
+        "headless": args.headless,
+        "width": "1920",
+        "height": "1080",
+    }
+)
+# Third Party
+from omni.isaac.core.utils.extensions import enable_extension
+
+ext_list = [
+    "omni.kit.asset_converter",
+    # "omni.kit.livestream.native",
+    "omni.kit.tool.asset_importer",
+    "omni.isaac.asset_browser",
+]
+[enable_extension(x) for x in ext_list]
+
+
+# Third Party
+import carb
+import numpy as np
+from helper import add_robot_to_scene
+from omni.isaac.core import World
+from omni.isaac.core.materials import OmniPBR
+from omni.isaac.core.objects import cuboid, sphere
+
+########### OV #################
+from omni.isaac.core.utils.extensions import enable_extension
+from omni.isaac.core.utils.types import ArticulationAction
+
+# CuRobo
+from curobo.util.logger import setup_curobo_logger
+from curobo.util.usd_helper import UsdHelper
+
+
+def main():
+    list_of_robots = get_motion_gen_robot_list()  # [:2]
+    usd_help = UsdHelper()
+
+    # assuming obstacles are in objects_path:
+    my_world = World(stage_units_in_meters=1.0)
+    my_world.scene.add_default_ground_plane()
+
+    stage = my_world.stage
+    usd_help.load_stage(stage)
+    xform = stage.DefinePrim("/World", "Xform")
+    stage.SetDefaultPrim(xform)
+    stage.DefinePrim("/curobo", "Xform")
+
+    # my_world.stage.SetDefaultPrim(my_world.stage.GetPrimAtPath("/World"))
+    stage = my_world.stage
+    # stage.SetDefaultPrim(stage.GetPrimAtPath("/World"))
+    # Make a target to follow
+    offset_y = 2
+    pose = Pose.from_list([0, 0, 0, 1, 0, 0, 0])
+
+    robot_cfg_list = []
+    robot_list = []
+    tensor_args = TensorDeviceType()
+    spheres = []
+
+    for i in range(len(list_of_robots)):
+        if i > 0:
+            pose.position[0, 1] += offset_y
+        if i == int(len(list_of_robots) / 2):
+            pose.position[0, 0] = -offset_y
+            pose.position[0, 1] = 0
+        robot_cfg = load_yaml(join_path(get_robot_configs_path(), list_of_robots[i]))["robot_cfg"]
+        robot_cfg_list.append(robot_cfg)
+        r = add_robot_to_scene(
+            robot_cfg,
+            my_world,
+            "/World/world_" + str(i) + "/",
+            robot_name="/World/world_" + str(i) + "/" "robot_" + str(i),
+            position=pose.position[0].cpu().numpy(),
+        )
+
+        robot_list.append(r[0])
+        robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+
+        kin_model = CudaRobotModel(robot_cfg.kinematics)
+        default_config = kin_model.cspace.retract_config
+
+        sph_list = kin_model.get_robot_as_spheres(default_config)
+        for si, s in enumerate(sph_list[0]):
+            sp = sphere.VisualSphere(
+                prim_path="/curobo/robot_sphere_" + str(i) + "_" + str(si),
+                position=np.ravel(s.position)
+                + pose.position[0].cpu().numpy()
+                + np.ravel([0, 0.5, 0.0]),
+                radius=float(s.radius),
+                color=np.array([0, 0.8, 0.2]),
+            )
+            spheres.append(sp)
+
+    setup_curobo_logger("warn")
+
+    while simulation_app.is_running():
+        my_world.step(render=True)
+
+        if not my_world.is_playing():
+            if i % 100 == 0:
+                print("**** Click Play to start simulation *****")
+            i += 1
+            continue
+        step_index = my_world.current_time_step_index
+
+        if step_index == 0:
+            my_world.reset()
+            for ri, robot in enumerate(robot_list):
+                j_names = robot_cfg_list[ri]["kinematics"]["cspace"]["joint_names"]
+                default_config = robot_cfg_list[ri]["kinematics"]["cspace"]["retract_config"]
+
+                idx_list = [robot.get_dof_index(x) for x in j_names]
+                robot.set_joint_positions(default_config, idx_list)
+                robot._articulation_view.set_max_efforts(
+                    values=np.array([5000 for i in range(len(idx_list))]), joint_indices=idx_list
+                )
+
+
+if __name__ == "__main__":
+    main()