158 lines
4.9 KiB
Python
158 lines
4.9 KiB
Python
#
|
|
# 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
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument(
|
|
"--headless_mode",
|
|
type=str,
|
|
default=None,
|
|
help="To run headless, use one of [native, websocket], webrtc might not work.",
|
|
)
|
|
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_mode is not None,
|
|
"width": "1920",
|
|
"height": "1080",
|
|
}
|
|
)
|
|
|
|
# Third Party
|
|
import carb
|
|
import numpy as np
|
|
from helper import add_extensions, add_robot_to_scene
|
|
from omni.isaac.core import World
|
|
from omni.isaac.core.objects import sphere
|
|
|
|
# CuRobo
|
|
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
|
|
|
|
# from curobo.wrap.reacher.ik_solver import IKSolver, IKSolverConfig
|
|
from curobo.types.base import TensorDeviceType
|
|
from curobo.types.math import Pose
|
|
from curobo.types.robot import RobotConfig
|
|
|
|
########### OV #################
|
|
from curobo.util.logger import setup_curobo_logger
|
|
from curobo.util.usd_helper import UsdHelper
|
|
from curobo.util_file import get_motion_gen_robot_list, get_robot_configs_path, join_path, load_yaml
|
|
|
|
|
|
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")
|
|
|
|
add_extensions(simulation_app, args.headless_mode)
|
|
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 <= 2:
|
|
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()
|