Add planning to grasp API

tests/goal_test.py

@@ -53,6 +53,7 @@ def test_repeat_seeds():
     weight = tensor_args.to_device([1, 1, 1, 1])
     vec_convergence = tensor_args.to_device([0, 0])
     run_weight = tensor_args.to_device([1])
+    project_distance = torch.tensor([True], device=tensor_args.device, dtype=torch.uint8)
     r = get_pose_distance(
         out_d,
         out_d.clone(),
@@ -72,6 +73,7 @@ def test_repeat_seeds():
         offset_waypoint,
         offset_tstep_fraction,
         g.batch_pose_idx,
+        project_distance,
         start_pose.position.shape[0],
         1,
         1,
@@ -79,7 +81,6 @@ def test_repeat_seeds():
         False,
         False,
         True,
-        True,
     )
 
     assert torch.sum(r[0]).item() <= 1e-5
@@ -105,6 +106,8 @@ def test_horizon_repeat_seeds():
     batch_pose_idx = torch.arange(0, b, 1, device=tensor_args.device, dtype=torch.int32).unsqueeze(
         -1
     )
+    project_distance = torch.tensor([True], device=tensor_args.device, dtype=torch.uint8)
+
     goal = Goal(goal_pose=goal_pose, batch_pose_idx=batch_pose_idx, current_state=current_state)
     g = goal  # .repeat_seeds(4)
 
@@ -144,6 +147,7 @@ def test_horizon_repeat_seeds():
         offset_waypoint,
         offset_tstep_fraction,
         g.batch_pose_idx,
+        project_distance,
         start_pose.position.shape[0],
         h,
         1,
@@ -151,6 +155,5 @@ def test_horizon_repeat_seeds():
         True,
         False,
         False,
-        True,
     )
     assert torch.sum(r[0]).item() < 1e-5

tests/motion_gen_goalset_test.py

@@ -20,11 +20,11 @@ from curobo.types.robot import JointState
 from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGenPlanConfig
 
 
-@pytest.fixture(scope="function")
+@pytest.fixture(scope="module")
 def motion_gen():
     tensor_args = TensorDeviceType()
     world_file = "collision_table.yml"
-    robot_file = "franka.yml"
+    robot_file = "ur5e_robotiq_2f_140.yml"
     motion_gen_config = MotionGenConfig.load_from_robot_config(
         robot_file,
         world_file,
@@ -202,3 +202,40 @@ def test_batch_goalset_padded(motion_gen_batch):
     q = result.optimized_plan.trim_trajectory(-1).squeeze(1)
     reached_state = motion_gen.compute_kinematics(q)
     assert torch.norm(goal_pose.position - reached_state.ee_pos_seq) < 0.005
+
+
+def test_grasp_goalset(motion_gen):
+    motion_gen.reset()
+    m_config = MotionGenPlanConfig(False, True)
+
+    retract_cfg = motion_gen.get_retract_config()
+
+    state = motion_gen.compute_kinematics(JointState.from_position(retract_cfg.view(1, -1)))
+
+    goal_pose = Pose(
+        state.ee_pos_seq.repeat(10, 1).view(1, -1, 3),
+        quaternion=state.ee_quat_seq.repeat(10, 1).view(1, -1, 4),
+    )
+    goal_pose.position[0, 0, 0] += 0.2
+
+    start_state = JointState.from_position(retract_cfg.view(1, -1) + 0.3)
+
+    result = motion_gen.plan_grasp(
+        start_state,
+        goal_pose,
+        m_config.clone(),
+        disable_collision_links=[
+            "left_outer_knuckle",
+            "left_inner_knuckle",
+            "left_outer_finger",
+            "left_inner_finger",
+            "left_inner_finger_pad",
+            "right_outer_knuckle",
+            "right_inner_knuckle",
+            "right_outer_finger",
+            "right_inner_finger",
+            "right_inner_finger_pad",
+        ],
+    )
+
+    assert torch.count_nonzero(result.success) == 1

tests/self_collision_test.py

@@ -9,6 +9,9 @@
 # its affiliates is strictly prohibited.
 #
 
+# Standard Library
+import copy
+
 # Third Party
 import pytest
 import torch
@@ -108,3 +111,73 @@ def test_self_collision_franka():
     cost_fn._out_distance[:] = 0.0
     out = cost_fn.forward(in_spheres)
     assert out.sum().item() > 0.0
+
+
+def test_self_collision_10k_spheres_franka():
+    tensor_args = TensorDeviceType()
+
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    robot_cfg["kinematics"]["debug"] = {"self_collision_experimental": False}
+
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    robot_cfg.kinematics.self_collision_config.experimental_kernel = True
+    kinematics = CudaRobotModel(robot_cfg.kinematics)
+    self_collision_data = kinematics.get_self_collision_config()
+    self_collision_config = SelfCollisionCostConfig(
+        **{"weight": 1.0, "classify": False, "self_collision_kin_config": self_collision_data},
+        tensor_args=tensor_args
+    )
+    cost_fn = SelfCollisionCost(self_collision_config)
+    cost_fn.self_collision_kin_config.experimental_kernel = True
+
+    b = 10
+    h = 1
+
+    q = torch.rand(
+        (b * h, kinematics.get_dof()), device=tensor_args.device, dtype=tensor_args.dtype
+    )
+
+    test_q = tensor_args.to_device([2.7735, -1.6737, 0.4998, -2.9865, 0.3386, 0.8413, 0.4371])
+    q[0, :] = test_q
+    kin_state = kinematics.get_state(q)
+
+    in_spheres = kin_state.link_spheres_tensor
+    in_spheres = in_spheres.view(b, h, -1, 4).contiguous()
+
+    out = cost_fn.forward(in_spheres)
+    assert out.sum().item() > 0.0
+
+    # create a franka robot with 10k spheres:
+    tensor_args = TensorDeviceType()
+
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    robot_cfg["kinematics"]["debug"] = {"self_collision_experimental": False}
+
+    sphere_cfg = load_yaml(
+        join_path(get_robot_configs_path(), robot_cfg["kinematics"]["collision_spheres"])
+    )["collision_spheres"]
+    n_times = 10
+    for k in sphere_cfg.keys():
+        sphere_cfg[k] = [copy.deepcopy(x) for x in sphere_cfg[k] for _ in range(n_times)]
+
+    robot_cfg["kinematics"]["collision_spheres"] = sphere_cfg
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    robot_cfg.kinematics.self_collision_config.experimental_kernel = False
+
+    kinematics = CudaRobotModel(robot_cfg.kinematics)
+    self_collision_data = kinematics.get_self_collision_config()
+    self_collision_config = SelfCollisionCostConfig(
+        **{"weight": 1.0, "classify": False, "self_collision_kin_config": self_collision_data},
+        tensor_args=tensor_args
+    )
+    cost_fn = SelfCollisionCost(self_collision_config)
+    cost_fn.self_collision_kin_config.experimental_kernel = False
+
+    kin_state = kinematics.get_state(q)
+
+    in_spheres = kin_state.link_spheres_tensor
+    in_spheres = in_spheres.view(b, h, -1, 4).contiguous()
+
+    out_10k = cost_fn.forward(in_spheres)
+    assert out_10k.sum().item() > 0.0
+    assert torch.linalg.norm(out - out_10k) < 1e-3

tests/voxelization_test.py

@@ -122,7 +122,7 @@ def test_esdf_from_world(world_collision):
     world_collision.clear_voxelization_cache()
     esdf = world_collision.get_esdf_in_bounding_box(voxel_size=voxel_size).clone()
 
-    occupied = esdf.get_occupied_voxels()
+    occupied = esdf.get_occupied_voxels(feature_threshold=0.0)
 
     assert voxels.shape == occupied.shape
 
@@ -136,7 +136,7 @@ def test_esdf_from_world(world_collision):
     indirect=True,
 )
 def test_voxels_prim_mesh(world_collision, world_collision_primitive):
-    voxel_size = 0.1
+    voxel_size = 0.05
     voxels = world_collision.get_voxels_in_bounding_box(voxel_size=voxel_size).clone()
     voxels_prim = world_collision_primitive.get_voxels_in_bounding_box(
         voxel_size=voxel_size
@@ -172,7 +172,7 @@ def test_esdf_prim_mesh(world_collision, world_collision_primitive):
     indirect=True,
 )
 def test_marching_cubes_from_world(world_collision):
-    voxel_size = 0.1
+    voxel_size = 0.05
     voxels = world_collision.get_voxels_in_bounding_box(voxel_size=voxel_size)
     mesh = Mesh.from_pointcloud(voxels[:, :3].detach().cpu().numpy(), pitch=voxel_size * 0.1)
 

tests/xrdf_test.py

@@ -12,6 +12,9 @@
 from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel, CudaRobotModelConfig
 from curobo.cuda_robot_model.util import load_robot_yaml
 from curobo.types.file_path import ContentPath
+from curobo.types.math import Pose
+from curobo.types.robot import JointState
+from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig
 
 
 def test_xrdf_kinematics():
@@ -31,3 +34,32 @@ def test_xrdf_kinematics():
 
     error = kin_pose.ee_pose.position - expected_position
     assert error.norm() < 0.01
+
+    assert "link_names" not in robot_data["robot_cfg"]["kinematics"]
+
+
+def test_xrdf_motion_gen():
+    robot_file = "ur10e.xrdf"
+    urdf_file = "robot/ur_description/ur10e.urdf"
+    content_path = ContentPath(robot_xrdf_file=robot_file, robot_urdf_file=urdf_file)
+    robot_data = load_robot_yaml(content_path)
+    robot_data["robot_cfg"]["kinematics"]["ee_link"] = "wrist_3_link"
+
+    motion_gen_config = MotionGenConfig.load_from_robot_config(
+        robot_data,
+        "collision_table.yml",
+        use_cuda_graph=True,
+        ee_link_name="tool0",
+    )
+    motion_gen = MotionGen(motion_gen_config)
+    motion_gen.warmup(warmup_js_trajopt=False)
+    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) + 0.3)
+    result = motion_gen.plan_single(start_state, retract_pose)
+
+    assert result.success.item()