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

tests/voxelization_test.py

@@ -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