release repository

tests/self_collision_test.py

@@ -0,0 +1,106 @@
+#
+# 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
+import torch
+
+# CuRobo
+from curobo.cuda_robot_model.cuda_robot_model import CudaRobotModel
+from curobo.rollout.cost.self_collision_cost import SelfCollisionCost, SelfCollisionCostConfig
+from curobo.types.base import TensorDeviceType
+from curobo.types.robot import RobotConfig
+from curobo.util_file import get_robot_configs_path, join_path, load_yaml
+
+
+@pytest.mark.parametrize(
+    "batch_size, horizon",
+    [
+        pytest.param(1, 1, id="1"),
+        pytest.param(10, 1, id="10"),
+        pytest.param(100000, 1, id="100k"),
+        pytest.param(100, 70, id="horizon"),
+    ],
+)
+def test_self_collision_experimental(batch_size, horizon):
+    robot_file = "franka.yml"
+    tensor_args = TensorDeviceType()
+
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), robot_file))["robot_cfg"]
+    robot_cfg["kinematics"]["debug"] = {"self_collision_experimental": True}
+    robot_cfg = RobotConfig.from_dict(robot_cfg, tensor_args)
+    kinematics = CudaRobotModel(robot_cfg.kinematics)
+    self_collision_data = kinematics.get_self_collision_config()
+    self_collision_config = SelfCollisionCostConfig(
+        **{"weight": 1.0, "classify": True, "self_collision_kin_config": self_collision_data},
+        tensor_args=tensor_args
+    )
+    cost_fn = SelfCollisionCost(self_collision_config)
+
+    b = batch_size
+    h = horizon
+
+    q = (
+        torch.rand(
+            (b * h, kinematics.get_dof()), device=tensor_args.device, dtype=tensor_args.dtype
+        )
+        * 10
+    )
+    kin_state = kinematics.get_state(q)
+
+    in_spheres = kin_state.link_spheres_tensor
+    in_spheres = in_spheres.view(b, h, -1, 4).contiguous()
+
+    for _ in range(1):
+        out = cost_fn.forward(in_spheres)
+    k = out.clone()
+    cost_fn.self_collision_kin_config.experimental_kernel = False
+    cost_fn._out_distance[:] = 0.0
+    for _ in range(1):
+        out = cost_fn.forward(in_spheres)
+
+    assert torch.norm(k - out).item() < 1e-8
+
+
+def test_self_collision_franka():
+    tensor_args = TensorDeviceType()
+
+    robot_cfg = load_yaml(join_path(get_robot_configs_path(), "franka.yml"))["robot_cfg"]
+    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": 5000.0, "classify": False, "self_collision_kin_config": self_collision_data},
+        tensor_args=tensor_args
+    )
+    cost_fn = SelfCollisionCost(self_collision_config)
+
+    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[:] = 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
+    cost_fn.self_collision_kin_config.experimental_kernel = False
+    cost_fn._out_distance[:] = 0.0
+    out = cost_fn.forward(in_spheres)
+    assert out.sum().item() > 0.0