hofee/gen_data_curobo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improved precision, quality and js planner.

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2024-04-11 13:19:01 -07:00
parent 774dcfd609
commit d6e600c88c
51 changed files with 2128 additions and 1054 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/conftest.py Normal file
View File

@@ -0,0 +1,15 @@
#
# 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
import os
os.environ["CUROBO_TORCH_COMPILE_DISABLE"] = str(1)

13
tests/motion_gen_api_test.py
View File

@@ -80,28 +80,25 @@ def test_motion_gen_lock_js_update():
kin_state = motion_gen_instance.compute_kinematics(start_state)
ee_pose = kin_state.ee_pose.clone()
# test motion gen:
plan_start = start_state.clone()
plan_start.position[..., :-2] += 0.1
result = motion_gen_instance.plan_single(plan_start, ee_pose)
result = motion_gen_instance.plan_single(plan_start, ee_pose.clone())
assert result.success.item()
lock_js = {"base_x": 1.0, "base_y": 0.0, "base_z": 0.0}
lock_js = {"base_x": 2.0, "base_y": 0.0, "base_z": 0.0}
motion_gen_instance.update_locked_joints(lock_js, robot_config)
kin_state_new = motion_gen_instance.compute_kinematics(start_state)
ee_pose_shift = kin_state_new.ee_pose.clone()
assert torch.norm(ee_pose.position[..., 0] - ee_pose_shift.position[..., 0]).item() == 1.0
assert 2 - torch.norm(ee_pose.position[..., 0] - ee_pose_shift.position[..., 0]).item() <= 1e-5
assert torch.norm(ee_pose.position[..., 1:] - ee_pose_shift.position[..., 1:]).item() == 0.0
# test motion gen with new lock state:
result = motion_gen_instance.plan_single(plan_start, ee_pose_shift)
result = motion_gen_instance.plan_single(plan_start, ee_pose_shift.clone())
assert result.success.item()
result = motion_gen_instance.plan_single(
plan_start, ee_pose, MotionGenPlanConfig(max_attempts=3)
plan_start, ee_pose.clone(), MotionGenPlanConfig(max_attempts=3)
)
assert result.success.item() == False

19
tests/motion_gen_goalset_test.py
View File

@@ -109,32 +109,35 @@ def test_batch_goalset_padded(motion_gen_batch):
# run goalset planning
motion_gen.reset()
retract_cfg = motion_gen.get_retract_config()
retract_cfg = motion_gen.get_retract_config().clone()
state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
goal_pose = Pose(
state.ee_pos_seq.repeat(3 * 3, 1).view(3, -1, 3).contiguous(),
quaternion=state.ee_quat_seq.repeat(3 * 3, 1).view(3, -1, 4).contiguous(),
)
).clone()
goal_pose.position[0, 1, 1] = 0.2
goal_pose.position[1, 0, 1] = 0.2
goal_pose.position[2, 1, 1] = 0.2
retract_cfg = motion_gen.get_retract_config().clone()
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3)
m_config = MotionGenPlanConfig(False, True, max_attempts=10, enable_graph_attempt=20)
result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config.clone())
m_config = MotionGenPlanConfig(enable_graph_attempt=100, max_attempts=2)
result = motion_gen.plan_batch_goalset(start_state, goal_pose.clone(), m_config.clone())
# get final solutions:
assert torch.count_nonzero(result.success) == result.success.shape[0]
reached_state = motion_gen.compute_kinematics(result.optimized_plan.trim_trajectory(-1))
reached_state = motion_gen.compute_kinematics(
result.optimized_plan.trim_trajectory(-1).squeeze(1)
)
#
goal_position = torch.cat(
[
goal_pose.position[x, result.goalset_index[x], :].unsqueeze(0)
goal_pose.position[x, result.goalset_index[x], :].clone().unsqueeze(0)
for x in range(len(result.goalset_index))
]
)
@@ -161,7 +164,7 @@ def test_batch_goalset_padded(motion_gen_batch):
start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.2).repeat_seeds(3)
result = motion_gen.plan_batch_goalset(start_state, goal_pose, m_config)
result = motion_gen.plan_batch_goalset(start_state, goal_pose.clone(), m_config)
# get final solutions:
assert torch.count_nonzero(result.success) == result.success.shape[0]
@@ -192,7 +195,7 @@ def test_batch_goalset_padded(motion_gen_batch):
goal_pose.position[1, 0] -= 0.1
result = motion_gen.plan_batch(start_state, goal_pose.clone(), m_config)
result = motion_gen.plan_batch(start_state, goal_pose, m_config)
assert torch.count_nonzero(result.success) == 3
# get final solutions:

44
tests/motion_gen_js_test.py Normal file
View File

@@ -0,0 +1,44 @@
#
# 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.types.robot import JointState
from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
def test_motion_gen_plan_js():
world_file = "collision_table.yml"
robot_file = "ur5e.yml"
motion_gen_config = MotionGenConfig.load_from_robot_config(
robot_file,
world_file,
trajopt_tsteps=32,
use_cuda_graph=False,
num_trajopt_seeds=4,
fixed_iters_trajopt=True,
evaluate_interpolated_trajectory=True,
)
motion_gen = MotionGen(motion_gen_config)
motion_gen.warmup(warmup_js_trajopt=True)
retract_cfg = motion_gen.get_retract_config()
start_state = JointState.from_position(retract_cfg.view(1, -1).clone())
goal_state = JointState.from_position(retract_cfg.view(1, -1).clone())
goal_state.position[:] = torch.as_tensor(
[1.000, -2.2000, 1.9000, -1.3830, -1.5700, 0.0000], device=motion_gen.tensor_args.device
)
result = motion_gen.plan_single_js(start_state, goal_state, MotionGenPlanConfig(max_attempts=1))
assert result.success.item()

41
tests/multi_pose_test.py
View File

@@ -42,11 +42,12 @@ def test_multi_pose_franka(b_size: int):
world_cfg,
rotation_threshold=0.05,
position_threshold=0.005,
num_seeds=30,
num_seeds=16,
self_collision_check=True,
self_collision_opt=True,
use_cuda_graph=True,
tensor_args=tensor_args,
regularization=False,
)
ik_solver = IKSolver(ik_config)
@@ -60,3 +61,41 @@ def test_multi_pose_franka(b_size: int):
assert (
torch.count_nonzero(success).item() / b_size >= 0.9
) # we check if atleast 90% are successful
@pytest.mark.parametrize(
"b_size",
[1, 10, 100],
)
def test_multi_pose_hand(b_size: int):
tensor_args = TensorDeviceType()
world_file = "collision_table.yml"
robot_file = "iiwa_allegro.yml"
robot_data = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
robot_cfg = RobotConfig.from_dict(robot_data)
world_cfg = WorldConfig.from_dict(load_yaml(join_path(get_world_configs_path(), world_file)))
ik_config = IKSolverConfig.load_from_robot_config(
robot_cfg,
world_cfg,
rotation_threshold=0.05,
position_threshold=0.005,
num_seeds=16,
use_cuda_graph=True,
tensor_args=tensor_args,
regularization=False,
)
ik_solver = IKSolver(ik_config)
q_sample = ik_solver.sample_configs(b_size)
kin_state = ik_solver.fk(q_sample)
link_poses = kin_state.link_pose
goal = Pose(kin_state.ee_position, kin_state.ee_quaternion).clone()
result = ik_solver.solve_batch(goal, link_poses=link_poses)
success = result.success
assert (
torch.count_nonzero(success).item() / b_size >= 0.9
) # we check if atleast 90% are successful