246 lines
8.3 KiB
Python
246 lines
8.3 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
|
|
|
|
# CuRobo
|
|
from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
|
|
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 JointState, RobotConfig
|
|
from curobo.util.logger import setup_curobo_logger
|
|
from curobo.util.usd_helper import UsdHelper
|
|
from curobo.util_file import (
|
|
get_assets_path,
|
|
get_robot_configs_path,
|
|
get_world_configs_path,
|
|
join_path,
|
|
load_yaml,
|
|
)
|
|
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
|
|
|
|
|
|
def save_curobo_world_to_usd():
|
|
world_file = "collision_table.yml"
|
|
world_cfg = WorldConfig.from_dict(
|
|
load_yaml(join_path(get_world_configs_path(), world_file))
|
|
).get_mesh_world(process=False)
|
|
usd_helper = UsdHelper()
|
|
usd_helper.create_stage()
|
|
|
|
usd_helper.add_obstacles_to_stage(world_cfg)
|
|
|
|
usd_helper.write_stage_to_file("test.usd")
|
|
|
|
|
|
def get_trajectory(robot_file="franka.yml", dt=1.0 / 60.0):
|
|
tensor_args = TensorDeviceType()
|
|
world_file = "collision_test.yml"
|
|
motion_gen_config = MotionGenConfig.load_from_robot_config(
|
|
robot_file,
|
|
world_file,
|
|
tensor_args,
|
|
trajopt_tsteps=24,
|
|
collision_checker_type=CollisionCheckerType.PRIMITIVE,
|
|
use_cuda_graph=True,
|
|
num_trajopt_seeds=2,
|
|
num_graph_seeds=2,
|
|
evaluate_interpolated_trajectory=True,
|
|
interpolation_dt=dt,
|
|
self_collision_check=True,
|
|
)
|
|
motion_gen = MotionGen(motion_gen_config)
|
|
motion_gen.warmup()
|
|
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
|
|
robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
|
|
retract_cfg = motion_gen.get_retract_config()
|
|
state = motion_gen.rollout_fn.compute_kinematics(
|
|
JointState.from_position(retract_cfg.view(1, -1))
|
|
)
|
|
|
|
retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
|
|
start_state = JointState.from_position(retract_cfg.view(1, -1).clone() + 0.5)
|
|
result = motion_gen.plan_single(start_state, retract_pose)
|
|
traj = result.get_interpolated_plan() # optimized plan
|
|
return traj
|
|
|
|
|
|
def save_curobo_robot_world_to_usd(robot_file="franka.yml"):
|
|
tensor_args = TensorDeviceType()
|
|
world_file = "collision_test.yml"
|
|
world_model = WorldConfig.from_dict(
|
|
load_yaml(join_path(get_world_configs_path(), world_file))
|
|
).get_obb_world()
|
|
dt = 1 / 60.0
|
|
|
|
q_traj = get_trajectory(robot_file, dt)
|
|
if q_traj is not None:
|
|
q_start = q_traj[0]
|
|
|
|
UsdHelper.write_trajectory_animation_with_robot_usd(
|
|
robot_file,
|
|
world_model,
|
|
q_start,
|
|
q_traj,
|
|
save_path="test.usda",
|
|
robot_color=[0.5, 0.5, 0.2, 1.0],
|
|
base_frame="/grid_world_1",
|
|
)
|
|
|
|
|
|
def save_curobo_robot_to_usd(robot_file="franka.yml"):
|
|
# print(robot_file)
|
|
tensor_args = TensorDeviceType()
|
|
robot_cfg_y = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
|
|
robot_cfg_y["kinematics"]["use_usd_kinematics"] = True
|
|
print(
|
|
len(robot_cfg_y["kinematics"]["cspace"]["null_space_weight"]),
|
|
len(robot_cfg_y["kinematics"]["cspace"]["retract_config"]),
|
|
len(robot_cfg_y["kinematics"]["cspace"]["joint_names"]),
|
|
)
|
|
# print(robot_cfg_y)
|
|
robot_cfg = RobotConfig.from_dict(robot_cfg_y, tensor_args)
|
|
start = JointState.from_position(robot_cfg.cspace.retract_config)
|
|
retract_cfg = robot_cfg.cspace.retract_config.clone()
|
|
retract_cfg[0] = 0.5
|
|
|
|
# print(retract_cfg)
|
|
kin_model = CudaRobotModel(robot_cfg.kinematics)
|
|
position = retract_cfg
|
|
q_traj = JointState.from_position(position.unsqueeze(0))
|
|
q_traj.joint_names = kin_model.joint_names
|
|
# print(q_traj.joint_names)
|
|
usd_helper = UsdHelper()
|
|
# usd_helper.create_stage(
|
|
# "test.usd", timesteps=q_traj.position.shape[0] + 1, dt=dt, interpolation_steps=10
|
|
# )
|
|
world_file = "collision_table.yml"
|
|
world_model = WorldConfig.from_dict(
|
|
load_yaml(join_path(get_world_configs_path(), world_file))
|
|
).get_obb_world()
|
|
|
|
# print(q_traj.position.shape)
|
|
# usd_helper.load_robot_usd(robot_cfg.kinematics.usd_path, js)
|
|
usd_helper.write_trajectory_animation_with_robot_usd(
|
|
{"robot_cfg": robot_cfg_y},
|
|
world_model,
|
|
start,
|
|
q_traj,
|
|
save_path="test.usd",
|
|
# robot_asset_prim_path="/robot"
|
|
)
|
|
|
|
# usd_helper.save()
|
|
# usd_helper.write_stage_to_file("test.usda")
|
|
|
|
|
|
def read_world_from_usd(file_path: str):
|
|
usd_helper = UsdHelper()
|
|
usd_helper.load_stage_from_file(file_path)
|
|
# world_model = usd_helper.get_obstacles_from_stage(reference_prim_path="/Root/world_obstacles")
|
|
world_model = usd_helper.get_obstacles_from_stage(
|
|
only_paths=["/world/obstacles"], reference_prim_path="/world"
|
|
)
|
|
# print(world_model)
|
|
for x in world_model.cuboid:
|
|
print(x.name + ":")
|
|
print(" pose: ", x.pose)
|
|
print(" dims: ", x.dims)
|
|
|
|
|
|
def read_robot_from_usd(robot_file: str = "franka.yml"):
|
|
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
|
|
robot_cfg["kinematics"]["use_usd_kinematics"] = True
|
|
robot_cfg = RobotConfig.from_dict(robot_cfg, TensorDeviceType())
|
|
|
|
|
|
def save_log_motion_gen(robot_file: str = "franka.yml"):
|
|
# load motion generation with debug mode:
|
|
robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
|
|
# robot_cfg["kinematics"]["collision_link_names"] = None
|
|
world_cfg = WorldConfig.from_dict(
|
|
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
|
|
).get_obb_world()
|
|
|
|
c_cache = {"obb": 10}
|
|
|
|
robot_cfg_instance = RobotConfig.from_dict(robot_cfg, tensor_args=TensorDeviceType())
|
|
|
|
enable_debug = True
|
|
motion_gen_config = MotionGenConfig.load_from_robot_config(
|
|
robot_cfg_instance,
|
|
world_cfg,
|
|
collision_cache=c_cache,
|
|
store_ik_debug=enable_debug,
|
|
store_trajopt_debug=enable_debug,
|
|
# num_ik_seeds=2,
|
|
# num_trajopt_seeds=1,
|
|
# ik_opt_iters=20,
|
|
# ik_particle_opt=False,
|
|
)
|
|
mg = MotionGen(motion_gen_config)
|
|
# mg.warmup(enable_graph=True, warmup_js_trajopt=False)
|
|
motion_gen = mg
|
|
# generate a plan:
|
|
retract_cfg = motion_gen.get_retract_config()
|
|
state = motion_gen.rollout_fn.compute_kinematics(
|
|
JointState.from_position(retract_cfg.view(1, -1))
|
|
)
|
|
link_chain = motion_gen.kinematics.kinematics_config.link_chain_map[
|
|
motion_gen.kinematics.kinematics_config.store_link_map.to(dtype=torch.long)
|
|
]
|
|
|
|
# exit()
|
|
link_poses = state.link_pose
|
|
# print(link_poses)
|
|
# del link_poses["tool0"]
|
|
# del link_poses["tool1"]
|
|
# del link_poses["tool2"]
|
|
|
|
# del link_poses["tool3"]
|
|
# print(link_poses)
|
|
|
|
retract_pose = Pose(state.ee_pos_seq.squeeze(), quaternion=state.ee_quat_seq.squeeze())
|
|
start_state = JointState.from_position(retract_cfg.view(1, -1).clone() + 0.5)
|
|
|
|
# get link poses if they exist:
|
|
|
|
result = motion_gen.plan_single(
|
|
start_state,
|
|
retract_pose,
|
|
link_poses=link_poses,
|
|
plan_config=MotionGenPlanConfig(max_attempts=1, partial_ik_opt=False),
|
|
)
|
|
UsdHelper.write_motion_gen_log(
|
|
result,
|
|
robot_cfg,
|
|
world_cfg,
|
|
start_state,
|
|
retract_pose,
|
|
join_path("log/usd/", "debug"),
|
|
write_robot_usd_path=join_path("log/usd/", "debug/assets/"),
|
|
write_ik=True,
|
|
write_trajopt=True,
|
|
visualize_robot_spheres=False,
|
|
grid_space=2,
|
|
link_poses=link_poses,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# save_curobo_world_to_usd()
|
|
setup_curobo_logger("error")
|
|
save_log_motion_gen("franka.yml")
|
|
# save_curobo_robot_world_to_usd("franka.yml")
|