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

Significantly improved convergence for mesh and cuboid, new ESDF collision.

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2024-03-18 11:19:48 -07:00
parent 286b3820a5
commit b1f63e8778
100 changed files with 7587 additions and 2589 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/cost_test.py
View File

@@ -41,7 +41,7 @@ def test_primitive_collision_cost():
)
cost = PrimitiveCollisionCost(cost_cfg)
q_spheres = torch.as_tensor(
[[0.1, 0.0, 0.0, 0.2], [10.0, 0.0, 0.0, 0.1]], **vars(tensor_args)
[[0.1, 0.0, 0.0, 0.2], [10.0, 0.0, 0.0, 0.1]], **(tensor_args.as_torch_dict())
).view(-1, 1, 1, 4)
c = cost.forward(q_spheres).flatten()
assert c[0] > 0.0 and c[1] == 0.0

18
tests/geom_test.py
View File

@@ -33,7 +33,8 @@ def test_world_primitive():
coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
coll_check = WorldPrimitiveCollision(coll_cfg)
x_sph = torch.as_tensor(
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(1, 1, -1, 4)
# create buffers:
query_buffer = CollisionQueryBuffer.initialize_from_shape(
@@ -58,7 +59,8 @@ def test_batch_world_primitive():
coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
coll_check = WorldPrimitiveCollision(coll_cfg)
x_sph = torch.as_tensor(
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(-1, 1, 1, 4)
env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
# create buffers:
@@ -90,7 +92,8 @@ def test_swept_world_primitive():
new_cube = Cuboid("cube_1", [0, 0, 1, 1, 0, 0, 0], None, [0.1, 0.2, 0.2])
coll_check.add_obb(new_cube, 0)
x_sph = torch.as_tensor(
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(1, 1, -1, 4)
env_query_idx = None
env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
@@ -125,7 +128,8 @@ def test_world_primitive_mesh_instance():
coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
coll_check = WorldMeshCollision(coll_cfg)
x_sph = torch.as_tensor(
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(1, 1, -1, 4)
# create buffers:
query_buffer = CollisionQueryBuffer.initialize_from_shape(
@@ -150,7 +154,8 @@ def test_batch_world_primitive_mesh_instance():
coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
coll_check = WorldMeshCollision(coll_cfg)
x_sph = torch.as_tensor(
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(-1, 1, 1, 4)
env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)
# create buffers:
@@ -184,7 +189,8 @@ def test_swept_world_primitive_mesh_instance():
w_obj_pose = Pose.from_list([0, 0, 1, 1, 0, 0, 0], tensor_args)
coll_check.add_obb_from_raw("cube_1", dims, 0, w_obj_pose=w_obj_pose)
x_sph = torch.as_tensor(
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(1, 1, -1, 4)
env_query_idx = None
env_query_idx = torch.zeros((x_sph.shape[0]), device=tensor_args.device, dtype=torch.int32)

10
tests/interpolation_test.py
View File

@@ -19,10 +19,11 @@ from curobo.util.trajectory import InterpolateType, get_batch_interpolated_traje
def test_linear_interpolation():
b, h, dof = 1, 24, 1
raw_dt = 0.4
int_dt = 0.01
tensor_args = TensorDeviceType()
b, h, dof = 1, 24, 1
raw_dt = tensor_args.to_device(0.4)
int_dt = 0.01
# initialize raw trajectory:
in_traj = JointState.zeros((b, h, dof), tensor_args)
in_traj.position = torch.zeros((b, h, dof), device=tensor_args.device)
@@ -62,6 +63,3 @@ def test_linear_interpolation():
).item()
< 0.05
)
# test_linear_interpolation()

30
tests/kinematics_test.py
View File

@@ -61,11 +61,15 @@ def test_franka_kinematics(cfg):
tensor_args = TensorDeviceType()
robot_model = CudaRobotModel(cfg)
q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
ee_position = torch.as_tensor([6.0860e-02, -4.7547e-12, 7.6373e-01], **vars(tensor_args)).view(
1, -1
)
ee_quat = torch.as_tensor([0.0382, 0.9193, 0.3808, 0.0922], **vars(tensor_args)).view(1, -1)
q_test = torch.as_tensor(
[0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **(tensor_args.as_torch_dict())
).view(1, -1)
ee_position = torch.as_tensor(
[6.0860e-02, -4.7547e-12, 7.6373e-01], **(tensor_args.as_torch_dict())
).view(1, -1)
ee_quat = torch.as_tensor(
[0.0382, 0.9193, 0.3808, 0.0922], **(tensor_args.as_torch_dict())
).view(1, -1)
b_list = [1, 10, 100, 5000][:1]
for b in b_list:
state = robot_model.get_state(q_test.repeat(b, 1).clone())
@@ -80,7 +84,9 @@ def test_franka_attached_object_kinematics(cfg):
tensor_args = TensorDeviceType()
robot_model = CudaRobotModel(cfg)
q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
q_test = torch.as_tensor(
[0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **(tensor_args.as_torch_dict())
).view(1, -1)
sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
state = robot_model.get_state(q_test.clone())
@@ -91,7 +97,7 @@ def test_franka_attached_object_kinematics(cfg):
assert torch.norm(attached_spheres[:, :, 1]) < 0.0001
# attach an object:
new_object = torch.zeros((2, 4), **vars(tensor_args))
new_object = torch.zeros((2, 4), **(tensor_args.as_torch_dict()))
new_object[:, 3] = 0.01
new_object[:, 1] = 0.1
robot_model.kinematics_config.update_link_spheres("attached_object", new_object)
@@ -106,7 +112,9 @@ def test_franka_attach_object_kinematics(cfg):
tensor_args = TensorDeviceType()
robot_model = CudaRobotModel(cfg)
q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
q_test = torch.as_tensor(
[0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **(tensor_args.as_torch_dict())
).view(1, -1)
sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("attached_object")
state = robot_model.get_state(q_test.clone())
@@ -117,7 +125,7 @@ def test_franka_attach_object_kinematics(cfg):
assert torch.norm(attached_spheres[:, :, 1]) < 0.0001
# attach an object:
new_object = torch.zeros((4, 4), **vars(tensor_args))
new_object = torch.zeros((4, 4), **(tensor_args.as_torch_dict()))
new_object[:2, 3] = 0.01
new_object[:2, 1] = 0.1
# print(attached_spheres[:, sph_idx,:].shape)
@@ -175,7 +183,9 @@ def test_franka_toggle_link_collision(cfg):
tensor_args = TensorDeviceType()
robot_model = CudaRobotModel(cfg)
q_test = torch.as_tensor([0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **vars(tensor_args)).view(1, -1)
q_test = torch.as_tensor(
[0.0, -1.2, 0.0, -2.0, 0.0, 1.0, 0.0], **(tensor_args.as_torch_dict())
).view(1, -1)
sph_idx = robot_model.kinematics_config.get_sphere_index_from_link_name("panda_link5")
state = robot_model.get_state(q_test.clone())

2
tests/motion_gen_module_test.py
View File

@@ -327,4 +327,4 @@ def test_motion_gen_single_js(motion_gen_str, enable_graph, request):
reached_state = result.optimized_plan[-1]
assert torch.norm(goal_state.position - reached_state.position) < 0.005
assert torch.norm(goal_state.position - reached_state.position) < 0.05

25
tests/nvblox_test.py
View File

@@ -17,9 +17,11 @@ import pytest
import torch
# CuRobo
from curobo.geom.sdf.world import CollisionQueryBuffer, WorldCollisionConfig
from curobo.geom.types import Cuboid, WorldConfig
from curobo.geom.sdf.world import CollisionCheckerType, CollisionQueryBuffer, WorldCollisionConfig
from curobo.geom.types import BloxMap, Cuboid, WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.util_file import get_world_configs_path, join_path, load_yaml
try:
@@ -41,7 +43,8 @@ def test_world_blox_trajectory():
coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
coll_check = WorldBloxCollision(coll_cfg)
x_sph = torch.as_tensor(
[[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(1, -1, 1, 4)
# create buffers:
query_buffer = CollisionQueryBuffer.initialize_from_shape(
@@ -66,7 +69,8 @@ def test_world_blox():
coll_cfg = WorldCollisionConfig(world_model=world_cfg, tensor_args=TensorDeviceType())
coll_check = WorldBloxCollision(coll_cfg)
x_sph = torch.as_tensor(
[[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.1, 0.2, 0.0, 10.5], [10.0, 0.0, 0.0, 0.0], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(1, 1, -1, 4)
# create buffers:
query_buffer = CollisionQueryBuffer.initialize_from_shape(
@@ -116,15 +120,12 @@ def test_nvblox_world_from_primitive_world():
data_dict = load_yaml(join_path(get_world_configs_path(), world_file))
world_cfg = WorldConfig.from_dict(data_dict).get_mesh_world(True)
mesh = world_cfg.mesh[0].get_trimesh_mesh()
# world_cfg.mesh[0].save_as_mesh("world.obj")
# Third Party
from nvblox_torch.datasets.mesh_dataset import MeshDataset
# CuRobo
from curobo.geom.sdf.world import CollisionCheckerType
from curobo.geom.types import BloxMap
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
try:
# Third Party
from nvblox_torch.datasets.mesh_dataset import MeshDataset
except ImportError:
pytest.skip("pyrender and scikit-image is not installed")
# create a nvblox collision checker:
world_config = WorldConfig(

4
tests/pose_reaching_test.py
View File

@@ -21,9 +21,7 @@ from curobo.wrap.reacher.motion_gen import MotionGen, MotionGenConfig, MotionGen
"parallel_finetune, force_graph, expected_motion_time",
[
(True, False, 12),
(False, False, 12),
(True, True, 12),
(False, True, 12),
],
)
def test_pose_sequence_ur5e(parallel_finetune, force_graph, expected_motion_time):
@@ -63,7 +61,7 @@ def test_pose_sequence_ur5e(parallel_finetune, force_graph, expected_motion_time
start_state.clone(),
goal_pose,
plan_config=MotionGenPlanConfig(
parallel_finetune=parallel_finetune, max_attempts=1, enable_graph=force_graph
parallel_finetune=parallel_finetune, max_attempts=10, enable_graph=force_graph
),
)
if result.success.item():

209
tests/robot_segmentation_test.py Normal file
View File

@@ -0,0 +1,209 @@
#
# 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.geom.types import PointCloud, WorldConfig
from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.types.robot import RobotConfig
from curobo.types.state import JointState
from curobo.util_file import get_robot_configs_path, get_world_configs_path, join_path, load_yaml
from curobo.wrap.model.robot_segmenter import RobotSegmenter
torch.manual_seed(30)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
try:
# Third Party
from nvblox_torch.datasets.mesh_dataset import MeshDataset
except ImportError:
pytest.skip(
"Nvblox Torch is not available or pyrender is not installed", allow_module_level=True
)
def create_render_dataset(
robot_file,
fov_deg: float = 60,
n_frames: int = 20,
retract_delta: float = 0.0,
):
# load robot:
robot_dict = load_yaml(join_path(get_robot_configs_path(), robot_file))
robot_dict["robot_cfg"]["kinematics"]["load_link_names_with_mesh"] = True
robot_dict["robot_cfg"]["kinematics"]["load_meshes"] = True
robot_cfg = RobotConfig.from_dict(robot_dict["robot_cfg"])
kin_model = CudaRobotModel(robot_cfg.kinematics)
q = kin_model.retract_config
q += retract_delta
meshes = kin_model.get_robot_as_mesh(q)
world = WorldConfig(mesh=meshes[:])
world_table = WorldConfig.from_dict(
load_yaml(join_path(get_world_configs_path(), "collision_table.yml"))
)
world_table.cuboid[0].dims = [0.5, 0.5, 0.1]
world.add_obstacle(world_table.objects[0])
robot_mesh = (
WorldConfig.create_merged_mesh_world(world, process_color=False).mesh[0].get_trimesh_mesh()
)
mesh_dataset = MeshDataset(
None,
n_frames=n_frames,
image_size=480,
save_data_dir=None,
trimesh_mesh=robot_mesh,
fov_deg=fov_deg,
)
q_js = JointState(position=q, joint_names=kin_model.joint_names)
return mesh_dataset, q_js
@pytest.mark.parametrize(
"robot_file",
["iiwa.yml", "iiwa_allegro.yml", "franka.yml", "ur10e.yml", "ur5e.yml"],
)
def test_mask_image(robot_file):
# create robot segmenter:
tensor_args = TensorDeviceType()
curobo_segmenter = RobotSegmenter.from_robot_file(
robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
)
mesh_dataset, q_js = create_render_dataset(robot_file, n_frames=5)
for i in range(len(mesh_dataset)):
data = mesh_dataset[i]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]).unsqueeze(0) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
if not curobo_segmenter.ready:
curobo_segmenter.update_camera_projection(cam_obs)
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs,
q_js,
)
if torch.count_nonzero(depth_mask) > 100:
return
assert False
@pytest.mark.parametrize(
"robot_file",
["iiwa.yml", "iiwa_allegro.yml", "franka.yml", "ur10e.yml", "ur5e.yml"],
)
def test_batch_mask_image(robot_file):
# create robot segmenter:
tensor_args = TensorDeviceType()
curobo_segmenter = RobotSegmenter.from_robot_file(
robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
)
mesh_dataset, q_js = create_render_dataset(robot_file, n_frames=5)
mesh_dataset_zoom, q_js = create_render_dataset(robot_file, fov_deg=40, n_frames=5)
for i in range(len(mesh_dataset)):
data = mesh_dataset[i]
data_zoom = mesh_dataset_zoom[i]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
cam_obs_zoom = CameraObservation(
depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
intrinsics=data_zoom["intrinsics"],
pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
)
cam_obs = cam_obs.stack(cam_obs_zoom)
if not curobo_segmenter.ready:
curobo_segmenter.update_camera_projection(cam_obs)
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs,
q_js,
)
if torch.count_nonzero(depth_mask[0]) > 100 and (torch.count_nonzero(depth_mask[1]) > 100):
return
assert False
@pytest.mark.parametrize(
"robot_file",
["iiwa.yml", "iiwa_allegro.yml", "franka.yml", "ur10e.yml", "ur5e.yml"],
)
def test_batch_robot_mask_image(robot_file):
# create robot segmenter:
tensor_args = TensorDeviceType()
curobo_segmenter = RobotSegmenter.from_robot_file(
robot_file, collision_sphere_buffer=0.01, distance_threshold=0.05, use_cuda_graph=True
)
mesh_dataset, q_js = create_render_dataset(robot_file, n_frames=5)
mesh_dataset_zoom, q_js_zoom = create_render_dataset(
robot_file,
fov_deg=40,
n_frames=5,
retract_delta=0.4,
)
q_js = q_js.unsqueeze(0).stack(q_js_zoom.unsqueeze(0))
for i in range(len(mesh_dataset)):
data = mesh_dataset[i]
data_zoom = mesh_dataset_zoom[i]
cam_obs = CameraObservation(
depth_image=tensor_args.to_device(data["depth"]) * 1000,
intrinsics=data["intrinsics"],
pose=Pose.from_matrix(data["pose"].to(device=tensor_args.device)),
)
cam_obs_zoom = CameraObservation(
depth_image=tensor_args.to_device(data_zoom["depth"]) * 1000,
intrinsics=data_zoom["intrinsics"],
pose=Pose.from_matrix(data_zoom["pose"].to(device=tensor_args.device)),
)
cam_obs = cam_obs.stack(cam_obs_zoom)
if not curobo_segmenter.ready:
curobo_segmenter.update_camera_projection(cam_obs)
depth_mask, filtered_image = curobo_segmenter.get_robot_mask_from_active_js(
cam_obs,
q_js,
)
if torch.count_nonzero(depth_mask[0]) > 100 and (torch.count_nonzero(depth_mask[1]) > 100):
return
assert False

254
tests/voxel_collision_test.py Normal file
View File

@@ -0,0 +1,254 @@
#
# 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.geom.sdf.world import (
CollisionCheckerType,
CollisionQueryBuffer,
WorldCollisionConfig,
WorldPrimitiveCollision,
)
from curobo.geom.sdf.world_mesh import WorldMeshCollision
from curobo.geom.sdf.world_voxel import WorldVoxelCollision
from curobo.geom.types import Cuboid, VoxelGrid, WorldConfig
from curobo.types.base import TensorDeviceType
def get_world_model(single_object: bool = False):
if single_object:
world_model = WorldConfig.from_dict(
{
"cuboid": {
"block1": {"dims": [0.1, 0.2, 0.5], "pose": [0.25, 0.1, 0, 1, 0, 0, 0]},
}
}
)
else:
world_model = WorldConfig.from_dict(
{
"cuboid": {
"block": {"dims": [0.5, 0.5, 0.5], "pose": [-0.25, 0, 0, 1, 0, 0, 0]},
"block1": {"dims": [0.1, 0.2, 0.5], "pose": [0.25, 0.1, 0, 1, 0, 0, 0]},
}
}
)
return world_model
@pytest.fixture(scope="function")
def world_collision(request):
world_model = get_world_model(request.param[1])
if request.param[0]:
world_model = world_model.get_mesh_world()
tensor_args = TensorDeviceType()
world_collision_config = WorldCollisionConfig.load_from_dict(
{
"checker_type": (
CollisionCheckerType.PRIMITIVE
if not request.param[0]
else CollisionCheckerType.MESH
),
"max_distance": 5.0,
"n_envs": 1,
},
world_model,
tensor_args,
)
if request.param[0]:
world_collision = WorldMeshCollision(world_collision_config)
else:
world_collision = WorldPrimitiveCollision(world_collision_config)
return world_collision
def world_voxel_collision_checker():
world_model = {
"voxel": {
"base": {"dims": [2.0, 2.0, 2.0], "pose": [0, 0, 0, 1, 0, 0, 0], "voxel_size": 0.05},
}
}
tensor_args = TensorDeviceType()
world_collision_config = WorldCollisionConfig.load_from_dict(
{
"checker_type": CollisionCheckerType.VOXEL,
"max_distance": 5.0,
"n_envs": 1,
},
world_model,
tensor_args,
)
world_collision = WorldVoxelCollision(world_collision_config)
return world_collision
@pytest.mark.parametrize(
"world_collision",
[
([True, True]),
([False, True]),
([True, False]),
([False, False]),
],
indirect=True,
)
def test_voxel_esdf(world_collision):
# create voxel collision checker
world_voxel_collision = world_voxel_collision_checker()
voxel_grid = world_voxel_collision.get_voxel_grid("base")
esdf = world_collision.get_esdf_in_bounding_box(
Cuboid(name="base", pose=voxel_grid.pose, dims=voxel_grid.dims),
voxel_size=voxel_grid.voxel_size,
)
world_voxel_collision.update_voxel_data(esdf)
voxel_size = 0.01
esdf = world_collision.get_esdf_in_bounding_box(
Cuboid(name="base", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1, 1, 1]), voxel_size=voxel_size
)
esdf_voxel = world_voxel_collision.get_esdf_in_bounding_box(
Cuboid(name="base", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1, 1, 1]), voxel_size=voxel_size
)
esdf_data = esdf.feature_tensor
esdf_voxel_data = esdf_voxel.feature_tensor
esdf_data[esdf_data < -1.0] = 0.0
esdf_voxel_data[esdf_voxel_data < -1.0] = 0.0
error = torch.abs(esdf_data - esdf_voxel_data)
assert torch.max(error) < 2 * voxel_grid.voxel_size
@pytest.mark.parametrize(
"world_collision",
[
([True, True]),
([False, True]),
([True, False]),
([False, False]),
],
indirect=True,
)
def test_primitive_voxel_sphere_distance(world_collision):
tensor_args = TensorDeviceType()
voxel_size = 0.025
world_voxel_collision = world_voxel_collision_checker()
voxel_grid = world_voxel_collision.get_voxel_grid("base")
esdf = world_collision.get_esdf_in_bounding_box(
Cuboid(name="base", pose=voxel_grid.pose, dims=voxel_grid.dims),
voxel_size=voxel_grid.voxel_size,
)
world_voxel_collision.update_voxel_data(esdf)
# create a grid and compute distance:
sample_grid = VoxelGrid(
name="test", pose=[0, 0, 0, 1, 0, 0, 0], voxel_size=voxel_size, dims=[1, 1, 1]
)
sample_spheres = sample_grid.create_xyzr_tensor()
sample_spheres = sample_spheres.reshape(-1, 1, 1, 4)
act_distance = tensor_args.to_device([0.0])
weight = tensor_args.to_device([1.0])
cuboid_collision_buffer = CollisionQueryBuffer.initialize_from_shape(
sample_spheres.shape, tensor_args, world_collision.collision_types
)
voxel_collision_buffer = CollisionQueryBuffer.initialize_from_shape(
sample_spheres.shape, tensor_args, world_voxel_collision.collision_types
)
d_cuboid = world_collision.get_sphere_distance(
sample_spheres, cuboid_collision_buffer, weight, act_distance
)
d_voxel = world_voxel_collision.get_sphere_distance(
sample_spheres, voxel_collision_buffer, weight, act_distance
)
error = torch.abs(d_cuboid.view(-1) - d_voxel.view(-1))
assert torch.max(error) < voxel_grid.voxel_size
@pytest.mark.parametrize(
"world_collision",
[
# ([True, True]),
([False, True]),
# ([True, False]),
# ([False, False]),
],
indirect=True,
)
def test_primitive_voxel_sphere_gradient(world_collision):
tensor_args = TensorDeviceType()
world_voxel_collision = world_voxel_collision_checker()
voxel_grid = world_voxel_collision.get_voxel_grid("base")
esdf = world_collision.get_esdf_in_bounding_box(
Cuboid(name="base", pose=voxel_grid.pose, dims=voxel_grid.dims),
voxel_size=voxel_grid.voxel_size,
)
voxel_size = voxel_grid.voxel_size
world_voxel_collision.update_voxel_data(esdf)
# create a grid and compute distance:
sample_grid = VoxelGrid(
name="test", pose=[0.0, 0.0, 0, 1, 0, 0, 0], voxel_size=voxel_size, dims=[0.1, 0.1, 0.1]
)
# sample_grid = VoxelGrid(
# name="test", pose=[0.2, 0.0, 0, 1, 0, 0, 0], voxel_size=voxel_size,
# dims=[0.1, 0.1, 0.1]
# )
sample_spheres = sample_grid.create_xyzr_tensor(transform_to_origin=True)
sample_spheres = sample_spheres.reshape(-1, 1, 1, 4)
act_distance = tensor_args.to_device([0.0])
weight = tensor_args.to_device([1.0])
cuboid_collision_buffer = CollisionQueryBuffer.initialize_from_shape(
sample_spheres.shape, tensor_args, world_collision.collision_types
)
voxel_collision_buffer = CollisionQueryBuffer.initialize_from_shape(
sample_spheres.shape, tensor_args, world_voxel_collision.collision_types
)
sample_spheres_1 = sample_spheres.clone()
sample_spheres_1.requires_grad = True
d_cuboid = world_collision.get_sphere_distance(
sample_spheres_1, cuboid_collision_buffer, weight, act_distance
)
sample_spheres_2 = sample_spheres.clone()
sample_spheres_2.requires_grad = True
d_voxel = world_voxel_collision.get_sphere_distance(
sample_spheres_2, voxel_collision_buffer, weight, act_distance
)
error = torch.abs(d_cuboid.view(-1) - d_voxel.view(-1))
assert torch.max(error) < voxel_grid.voxel_size
cuboid_gradient = cuboid_collision_buffer.get_gradient_buffer()
voxel_gradient = voxel_collision_buffer.get_gradient_buffer()
error = torch.linalg.norm(cuboid_gradient - voxel_gradient, dim=-1)
print(cuboid_gradient)
print(voxel_gradient)
assert torch.max(error) < voxel_grid.voxel_size

181
tests/voxelization_test.py Normal file
View File

@@ -0,0 +1,181 @@
#
# 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
# CuRobo
from curobo.geom.sdf.world import (
CollisionCheckerType,
WorldCollisionConfig,
WorldPrimitiveCollision,
)
from curobo.geom.sdf.world_mesh import WorldMeshCollision
from curobo.geom.types import Mesh, WorldConfig
from curobo.types.base import TensorDeviceType
def get_world_model(single_object: bool = False):
if single_object:
world_model = WorldConfig.from_dict(
{
"cuboid": {
"block1": {"dims": [0.1, 0.2, 0.5], "pose": [0.25, 0.1, 0, 1, 0, 0, 0]},
}
}
)
else:
world_model = WorldConfig.from_dict(
{
"cuboid": {
"block": {"dims": [0.5, 0.5, 0.5], "pose": [-0.25, 0, 0, 1, 0, 0, 0]},
"block1": {"dims": [0.1, 0.2, 0.5], "pose": [0.25, 0.1, 0, 1, 0, 0, 0]},
}
}
)
return world_model
@pytest.fixture(scope="function")
def world_collision(request):
world_model = get_world_model(request.param[1])
if request.param[0]:
world_model = world_model.get_mesh_world()
tensor_args = TensorDeviceType()
world_collision_config = WorldCollisionConfig.load_from_dict(
{
"checker_type": (
CollisionCheckerType.PRIMITIVE
if not request.param[0]
else CollisionCheckerType.MESH
),
"max_distance": 5.0,
"n_envs": 1,
},
world_model,
tensor_args,
)
if request.param[0]:
world_collision = WorldMeshCollision(world_collision_config)
else:
world_collision = WorldPrimitiveCollision(world_collision_config)
return world_collision
@pytest.fixture(scope="function")
def world_collision_primitive():
world_model = get_world_model(True)
tensor_args = TensorDeviceType()
world_collision_config = WorldCollisionConfig.load_from_dict(
{
"checker_type": (CollisionCheckerType.PRIMITIVE),
"max_distance": 5.0,
"n_envs": 1,
},
world_model,
tensor_args,
)
world_collision = WorldPrimitiveCollision(world_collision_config)
return world_collision
@pytest.mark.parametrize(
"world_collision",
[
(True, True),
(False, True),
],
indirect=True,
)
def test_voxels_from_world(world_collision):
voxel_size = 0.1
voxels = world_collision.get_voxels_in_bounding_box(voxel_size=voxel_size)
assert voxels.shape[0] > 4
@pytest.mark.skip(reason="Not ready yet.")
@pytest.mark.parametrize(
"world_collision",
[
(True, True),
(False, True),
],
indirect=True,
)
def test_esdf_from_world(world_collision):
voxel_size = 0.02
voxels = world_collision.get_voxels_in_bounding_box(voxel_size=voxel_size).clone()
world_collision.clear_voxelization_cache()
esdf = world_collision.get_esdf_in_bounding_box(voxel_size=voxel_size).clone()
occupied = esdf.get_occupied_voxels()
assert voxels.shape == occupied.shape
@pytest.mark.parametrize(
"world_collision",
[
(True, True),
(False, True),
],
indirect=True,
)
def test_voxels_prim_mesh(world_collision, world_collision_primitive):
voxel_size = 0.1
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
).clone()
assert voxels.shape == voxels_prim.shape
@pytest.mark.parametrize(
"world_collision",
[
(True, True),
# (False, True),
],
indirect=True,
)
def test_esdf_prim_mesh(world_collision, world_collision_primitive):
voxel_size = 0.1
esdf = world_collision.get_esdf_in_bounding_box(voxel_size=voxel_size).clone()
esdf_prim = world_collision_primitive.get_esdf_in_bounding_box(voxel_size=voxel_size).clone()
print(esdf.voxel_size, esdf_prim.voxel_size)
voxels = esdf.get_occupied_voxels()
voxels_prim = esdf_prim.get_occupied_voxels()
assert voxels.shape == voxels_prim.shape
@pytest.mark.parametrize(
"world_collision",
[
(True, True),
(False, True),
(True, False),
(False, False),
],
indirect=True,
)
def test_marching_cubes_from_world(world_collision):
voxel_size = 0.1
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)
mesh.save_as_mesh("test_" + str(len(voxels)) + ".stl")
assert len(mesh.vertices) > 100 # exact value is 240

6
tests/warp_mesh_test.py
View File

@@ -32,7 +32,7 @@ def test_sdf_pose():
new_mesh,
env_idx=0,
)
query_spheres = torch.zeros((1, 1, 2, 4), **vars(tensor_args))
query_spheres = torch.zeros((1, 1, 2, 4), **(tensor_args.as_torch_dict()))
query_spheres[..., 2] = 10.0
query_spheres[..., 1, :] = 0.0
query_spheres[..., 3] = 1.0
@@ -59,7 +59,7 @@ def test_swept_sdf_speed_pose():
new_mesh,
env_idx=0,
)
query_spheres = torch.zeros((1, 1, 2, 4), **vars(tensor_args))
query_spheres = torch.zeros((1, 1, 2, 4), **(tensor_args.as_torch_dict()))
query_spheres[..., 2] = 10.0
query_spheres[..., 1, :] = 0.0
query_spheres[..., 3] = 1.0
@@ -88,7 +88,7 @@ def test_swept_sdf_pose():
new_mesh,
env_idx=0,
)
query_spheres = torch.zeros((1, 1, 2, 4), **vars(tensor_args))
query_spheres = torch.zeros((1, 1, 2, 4), **(tensor_args.as_torch_dict()))
query_spheres[..., 2] = 10.0
query_spheres[..., 1, :] = 0.0
query_spheres[..., 3] = 1.0

3
tests/world_config_test.py
View File

@@ -91,7 +91,8 @@ def test_world_modify():
world_ccheck.update_obstacle_pose(name="cylinder_1", w_obj_pose=w_pose)
x_sph = torch.as_tensor(
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.1], [0.01, 0.01, 0.0, 0.1]], **vars(tensor_args)
[[0.0, 0.0, 0.0, 0.1], [10.0, 0.0, 0.0, 0.1], [0.01, 0.01, 0.0, 0.1]],
**(tensor_args.as_torch_dict())
).view(1, 1, -1, 4)
# create buffers:
query_buffer = CollisionQueryBuffer.initialize_from_shape(