137 lines
4.7 KiB
Python
137 lines
4.7 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 pytest
|
|
|
|
try:
|
|
# CuRobo
|
|
from curobo.util.usd_helper import UsdHelper
|
|
except ImportError:
|
|
pytest.skip("usd-core not found, skipping usd tests", allow_module_level=True)
|
|
|
|
# CuRobo
|
|
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_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
|
|
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
|
|
|
|
|
|
@pytest.mark.skip(reason="Takes 60+ seconds and is not a core functionality")
|
|
def test_write_motion_gen_log(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,
|
|
)
|
|
mg = MotionGen(motion_gen_config)
|
|
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_poses = state.link_pose
|
|
|
|
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,
|
|
)
|
|
assert True
|
|
|
|
|
|
def test_write_trajectory_usd(robot_file="franka.yml"):
|
|
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
|
|
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,
|
|
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)
|
|
q_traj = result.get_interpolated_plan() # optimized plan
|
|
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",
|
|
dt=result.interpolation_dt,
|
|
)
|
|
else:
|
|
assert False
|