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

Add planning to grasp API

Browse Source
This commit is contained in:
Balakumar Sundaralingam
2024-11-22 14:15:18 -08:00
parent 18e9ebd35f
commit 36ea382dab
38 changed files with 939 additions and 535 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/curobo/geom/sdf/world.py
View File

@@ -633,6 +633,7 @@ class WorldCollision(WorldCollisionConfig):
self,
cuboid: Cuboid = Cuboid(name="test", pose=[0, 0, 0, 1, 0, 0, 0], dims=[1, 1, 1]),
voxel_size: float = 0.02,
run_marching_cubes: bool = True,
) -> Mesh:
"""Get a mesh representation of the world obstacles based on occupancy in a bounding box.
@@ -642,19 +643,31 @@ class WorldCollision(WorldCollisionConfig):
Args:
cuboid: Bounding box to get the mesh representation.
voxel_size: Size of the voxels in meters.
run_marching_cubes: Runs marching cubes over occupied voxels to generate a mesh. If
set to False, then all occupied voxels are merged into a mesh and returned.
Returns:
Mesh representation of the world obstacles in the bounding box.
"""
voxels = self.get_voxels_in_bounding_box(cuboid, voxel_size)
# voxels = voxels.cpu().numpy()
# cuboids = [Cuboid(name="c_"+str(x), pose=[voxels[x,0],voxels[x,1],voxels[x,2], 1,0,0,0],
# dims=[voxel_size, voxel_size, voxel_size]) for x in range(voxels.shape[0])]
# mesh = WorldConfig(cuboid=cuboids).get_mesh_world(True).mesh[0]
mesh = Mesh.from_pointcloud(
voxels[:, :3].detach().cpu().numpy(),
pitch=voxel_size * 1.1,
)
voxels = voxels.cpu().numpy()
if run_marching_cubes:
mesh = Mesh.from_pointcloud(
voxels[:, :3].detach().cpu().numpy(),
pitch=voxel_size * 1.1,
)
else:
cuboids = [
Cuboid(
name="c_" + str(x),
pose=[voxels[x, 0], voxels[x, 1], voxels[x, 2], 1, 0, 0, 0],
dims=[voxel_size, voxel_size, voxel_size],
)
for x in range(voxels.shape[0])
]
mesh = WorldConfig(cuboid=cuboids).get_mesh_world(True).mesh[0]
return mesh
def get_obstacle_names(self, env_idx: int = 0) -> List[str]:

28
src/curobo/geom/sdf/world_voxel.py
View File

@@ -358,12 +358,21 @@ class WorldVoxelCollision(WorldMeshCollision):
env_idx: Environment index to update voxel grid in.
"""
obs_idx = self.get_voxel_idx(new_voxel.name, env_idx)
self._voxel_tensor_list[3][env_idx, obs_idx, :, :] = new_voxel.feature_tensor.view(
new_voxel.feature_tensor.shape[0], -1
).to(dtype=self._voxel_tensor_list[3].dtype)
self._voxel_tensor_list[0][env_idx, obs_idx, :3] = self.tensor_args.to_device(
new_voxel.dims
)
feature_tensor = new_voxel.feature_tensor.view(new_voxel.feature_tensor.shape[0], -1)
if (
feature_tensor.shape[0] != self._voxel_tensor_list[3][env_idx, obs_idx, :, :].shape[0]
or feature_tensor.shape[1]
!= self._voxel_tensor_list[3][env_idx, obs_idx, :, :].shape[1]
):
log_error(
"Feature tensor shape mismatch, existing shape: "
+ str(self._voxel_tensor_list[3][env_idx, obs_idx, :, :].shape)
+ " New shape: "
+ str(feature_tensor.shape)
)
self._voxel_tensor_list[3][env_idx, obs_idx, :, :].copy_(feature_tensor)
self._voxel_tensor_list[0][env_idx, obs_idx, :3].copy_(torch.as_tensor(new_voxel.dims))
self._voxel_tensor_list[0][env_idx, obs_idx, 3] = new_voxel.voxel_size
self._voxel_tensor_list[1][env_idx, obs_idx, :7] = (
Pose.from_list(new_voxel.pose, self.tensor_args).inverse().get_pose_vector()
@@ -876,14 +885,19 @@ class WorldVoxelCollision(WorldMeshCollision):
self._env_n_voxels[:] = 0
super().clear_cache()
def get_voxel_grid_shape(self, env_idx: int = 0, obs_idx: int = 0) -> torch.Size:
def get_voxel_grid_shape(
self, env_idx: int = 0, obs_idx: int = 0, name: Optional[str] = None
) -> torch.Size:
"""Get dimensions of the voxel grid.
Args:
env_idx: Environment index.
obs_idx: Obstacle index.
name: Name of obstacle. When provided, obs_idx is ignored.
Returns:
Shape of the voxel grid.
"""
if name is not None:
obs_idx = self.get_voxel_idx(name, env_idx)
return self._voxel_tensor_list[3][env_idx, obs_idx].shape

33
src/curobo/geom/types.py
View File

@@ -13,7 +13,6 @@
from __future__ import annotations
# Standard Library
import math
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
@@ -28,6 +27,7 @@ from curobo.geom.sphere_fit import SphereFitType, fit_spheres_to_mesh
from curobo.types.base import TensorDeviceType
from curobo.types.camera import CameraObservation
from curobo.types.math import Pose
from curobo.util.helpers import robust_floor
from curobo.util.logger import log_error, log_warn
from curobo.util_file import get_assets_path, join_path
@@ -723,12 +723,16 @@ class VoxelGrid(Obstacle):
"""Get shape of voxel grid."""
bounds = self.dims
grid_shape = [bounds[0], bounds[1], bounds[2]]
inv_voxel_size = 1.0 / self.voxel_size
grid_shape = [1 + robust_floor(x * inv_voxel_size) for x in grid_shape]
low = [-bounds[0] / 2, -bounds[1] / 2, -bounds[2] / 2]
high = [bounds[0] / 2, bounds[1] / 2, bounds[2] / 2]
grid_shape = [
1 + int(high[i] / self.voxel_size) - (int(low[i] / self.voxel_size))
for i in range(len(low))
]
return grid_shape, low, high
def create_xyzr_tensor(
@@ -745,10 +749,19 @@ class VoxelGrid(Obstacle):
"""
trange, low, high = self.get_grid_shape()
x = torch.linspace(low[0], high[0], trange[0], device=tensor_args.device)
y = torch.linspace(low[1], high[1], trange[1], device=tensor_args.device)
z = torch.linspace(low[2], high[2], trange[2], device=tensor_args.device)
inv_voxel_size = 1.0 / self.voxel_size
x = torch.linspace(1, trange[0], trange[0], device=tensor_args.device) - round(
(0.5 * self.dims[0]) * inv_voxel_size
)
y = torch.linspace(1, trange[1], trange[1], device=tensor_args.device) - round(
(0.5 * self.dims[1]) * inv_voxel_size
)
z = torch.linspace(1, trange[2], trange[2], device=tensor_args.device) - round(
(0.5 * self.dims[2]) * inv_voxel_size
)
x = x * self.voxel_size - 0.5 * self.voxel_size
y = y * self.voxel_size - 0.5 * self.voxel_size
z = z * self.voxel_size - 0.5 * self.voxel_size
w, l, h = x.shape[0], y.shape[0], z.shape[0]
xyz = (
torch.stack(torch.meshgrid(x, y, z, indexing="ij")).permute((1, 2, 3, 0)).reshape(-1, 3)
@@ -757,7 +770,7 @@ class VoxelGrid(Obstacle):
if transform_to_origin:
pose = Pose.from_list(self.pose, tensor_args=tensor_args)
xyz = pose.transform_points(xyz.contiguous())
r = torch.zeros_like(xyz[:, 0:1]) + (self.voxel_size * 0.5)
r = torch.zeros_like(xyz[:, 0:1])
xyzr = torch.cat([xyz, r], dim=1)
return xyzr