184 lines
6.6 KiB
Python
184 lines
6.6 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.
|
|
#
|
|
|
|
# Standard Library
|
|
import copy
|
|
|
|
# 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": False}
|
|
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)
|
|
cost_fn.self_collision_kin_config.experimental_kernel = True
|
|
|
|
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["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": 5000.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[:] = 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
|
|
|
|
|
|
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
|