Remove wp.Volume (#683)

* remove wp.Volume

* auto-format

* fix sensor

* fix kernel arguments

* fix small bugs

* fix indentation

* add box_project function

* use analytic gradient where available

* remove redundant aabb check

* fix test requirement

Author: quagla

mujoco_warp/_src/collision_sdf.py

@@ -27,6 +27,8 @@
 from .types import GeomType
 from .types import Model
 from .types import vec5
+from .types import vec8f
+from .types import vec8i
 from .util_misc import halton
 from .warp_util import event_scope
 
@@ -47,9 +49,12 @@ class AABB:
 
 @wp.struct
 class VolumeData:
-  volume_id: wp.uint64
   center: wp.vec3
   half_size: wp.vec3
+  oct_aabb: wp.array2d(dtype=wp.vec3)
+  oct_child: wp.array(dtype=vec8i)
+  oct_coeff: wp.array(dtype=vec8f)
+  valid: bool = False
 
 
 @wp.struct
@@ -68,20 +73,12 @@ class MeshData:
   valid: bool = False
 
 
-@wp.func
-def get_volume_data(volume_id: wp.uint64, center: wp.vec3, half_size: wp.vec3) -> VolumeData:
-  volume_data = VolumeData()
-  volume_data.volume_id = volume_id
-  volume_data.center = center
-  volume_data.half_size = half_size
-  return volume_data
-
-
 @wp.func
 def get_sdf_params(
   # Model:
-  volume_ids: wp.array(dtype=wp.uint64),
   oct_aabb: wp.array2d(dtype=wp.vec3),
+  oct_child: wp.array(dtype=vec8i),
+  oct_coeff: wp.array(dtype=vec8f),
   plugin: wp.array(dtype=int),
   plugin_attr: wp.array(dtype=wp.vec3f),
   # In:
@@ -92,17 +89,19 @@ def get_sdf_params(
 ) -> Tuple[wp.vec3, int, VolumeData, MeshData]:
   attributes = g_size
   plugin_index = -1
-  volume_data = get_volume_data(wp.uint64(0), wp.vec3(0.0), wp.vec3(0.0))
+  volume_data = VolumeData()
 
   if g_type == int(GeomType.SDF.value) and plugin_id != -1:
     attributes = plugin_attr[plugin_id]
     plugin_index = plugin[plugin_id]
 
   elif g_type == int(GeomType.SDF.value) and mesh_id != -1:
-    volume_id = volume_ids[mesh_id]
-    center = oct_aabb[mesh_id, 0]
-    half_size = oct_aabb[mesh_id, 1]
-    volume_data = get_volume_data(volume_id, center, half_size)
+    volume_data.center = oct_aabb[mesh_id, 0]
+    volume_data.half_size = oct_aabb[mesh_id, 1]
+    volume_data.oct_aabb = oct_aabb
+    volume_data.oct_child = oct_child
+    volume_data.oct_coeff = oct_coeff
+    volume_data.valid = True
 
   return attributes, plugin_index, volume_data, MeshData()
 
@@ -223,22 +222,73 @@ def user_sdf_grad(p: wp.vec3, attr: wp.vec3, sdf_type: int) -> wp.vec3:
 
 
 @wp.func
-def sample_volume_sdf(xyz: wp.vec3, volume_data: VolumeData) -> float:
-  center = volume_data.center
-  half_size = volume_data.half_size
+def find_oct(
+  oct_aabb: wp.array2d(dtype=wp.vec3), oct_child: wp.array(dtype=vec8i), p: wp.vec3, grad: bool
+) -> Tuple[int, Tuple[vec8f, vec8f, vec8f]]:
+  stack = int(0)
+  niter = int(100)
+  rx = vec8f(0.0)
+  ry = vec8f(0.0)
+  rz = vec8f(0.0)
+
+  while niter > 0:
+    niter -= 1
+    node = stack
+
+    if node == -1:
+      wp.printf("ERROR: Invalid node number\n")
+      return -1, (rx, ry, rz)
+
+    vmin = oct_aabb[node, 0] - oct_aabb[node, 1]
+    vmax = oct_aabb[node, 0] + oct_aabb[node, 1]
+    coord = wp.cw_div(p - vmin, vmax - vmin)
+
+    # check if the node is a leaf
+    if (
+      oct_child[node][0] == -1
+      and oct_child[node][1] == -1
+      and oct_child[node][2] == -1
+      and oct_child[node][3] == -1
+      and oct_child[node][4] == -1
+      and oct_child[node][5] == -1
+      and oct_child[node][6] == -1
+      and oct_child[node][7] == -1
+    ):
+      for j in range(8):
+        if not grad:
+          rx[j] = (
+            (coord[0] if j & 1 else 1.0 - coord[0])
+            * (coord[1] if j & 2 else 1.0 - coord[1])
+            * (coord[2] if j & 4 else 1.0 - coord[2])
+          )
+        else:
+          rx[j] = (1.0 if j & 1 else -1.0) * (coord[1] if j & 2 else 1.0 - coord[1]) * (coord[2] if j & 4 else 1.0 - coord[2])
+          ry[j] = (coord[0] if j & 1 else 1.0 - coord[0]) * (1.0 if j & 2 else -1.0) * (coord[2] if j & 4 else 1.0 - coord[2])
+          rz[j] = (coord[0] if j & 1 else 1.0 - coord[0]) * (coord[1] if j & 2 else 1.0 - coord[1]) * (1.0 if j & 4 else -1.0)
+      return node, (rx, ry, rz)
+
+    # compute which of 8 children to visit next
+    x = 1 if coord[0] < 0.5 else 0
+    y = 1 if coord[1] < 0.5 else 0
+    z = 1 if coord[2] < 0.5 else 0
+    stack = oct_child[node][4 * z + 2 * y + x]
+
+  wp.print("ERROR: Node not found\n")
+  return -1, (rx, ry, rz)
 
+
+@wp.func
+def box_project(center: wp.vec3, half_size: wp.vec3, xyz: wp.vec3) -> Tuple[float, wp.vec3]:
   r = xyz - center
   q = wp.vec3(wp.abs(r[0]) - half_size[0], wp.abs(r[1]) - half_size[1], wp.abs(r[2]) - half_size[2])
 
   if q[0] <= 0.0 and q[1] <= 0.0 and q[2] <= 0.0:
-    uvw = wp.volume_world_to_index(volume_data.volume_id, xyz)
-    sdf = wp.volume_sample_f(volume_data.volume_id, uvw, wp.Volume.LINEAR)
-    return sdf
+    return 0.0, xyz
 
   else:
-    point = wp.vec3(xyz[0], xyz[1], xyz[2])
     dist_sqr = 0.0
     eps = 1e-4
+    point = wp.vec3(xyz[0], xyz[1], xyz[2])
 
     if q[0] >= 0.0:
       dist_sqr += q[0] * q[0]
@@ -261,23 +311,33 @@ def sample_volume_sdf(xyz: wp.vec3, volume_data: VolumeData) -> float:
       else:
         point = wp.vec3(point[0], point[1], point[2] + (q[2] + eps))
 
-    dist0 = wp.sqrt(dist_sqr)
+    return wp.sqrt(dist_sqr), point
 
-    uvw = wp.volume_world_to_index(volume_data.volume_id, point)
-    sdf = wp.volume_sample_f(volume_data.volume_id, uvw, wp.Volume.LINEAR)
-    return dist0 + sdf
+
+@wp.func
+def sample_volume_sdf(xyz: wp.vec3, volume_data: VolumeData) -> float:
+  dist0, point = box_project(volume_data.center, volume_data.half_size, xyz)
+  node, weights = find_oct(volume_data.oct_aabb, volume_data.oct_child, point, grad=False)
+  return dist0 + wp.dot(weights[0], volume_data.oct_coeff[node])
 
 
 @wp.func
 def sample_volume_grad(xyz: wp.vec3, volume_data: VolumeData) -> wp.vec3:
-  h = 1e-4
-  dx = wp.vec3(h, 0.0, 0.0)
-  dy = wp.vec3(0.0, h, 0.0)
-  dz = wp.vec3(0.0, 0.0, h)
-  f = sample_volume_sdf(xyz, volume_data)
-  grad_x = (sample_volume_sdf(xyz + dx, volume_data) - f) / h
-  grad_y = (sample_volume_sdf(xyz + dy, volume_data) - f) / h
-  grad_z = (sample_volume_sdf(xyz + dz, volume_data) - f) / h
+  dist0, point = box_project(volume_data.center, volume_data.half_size, xyz)
+  if dist0 > 0:
+    h = 1e-4
+    dx = wp.vec3(h, 0.0, 0.0)
+    dy = wp.vec3(0.0, h, 0.0)
+    dz = wp.vec3(0.0, 0.0, h)
+    f = sample_volume_sdf(xyz, volume_data)
+    grad_x = (sample_volume_sdf(xyz + dx, volume_data) - f) / h
+    grad_y = (sample_volume_sdf(xyz + dy, volume_data) - f) / h
+    grad_z = (sample_volume_sdf(xyz + dz, volume_data) - f) / h
+    return wp.vec3(grad_x, grad_y, grad_z)
+  node, weights = find_oct(volume_data.oct_aabb, volume_data.oct_child, point, grad=True)
+  grad_x = wp.dot(weights[0], volume_data.oct_coeff[node])
+  grad_y = wp.dot(weights[1], volume_data.oct_coeff[node])
+  grad_z = wp.dot(weights[2], volume_data.oct_coeff[node])
   return wp.vec3(grad_x, grad_y, grad_z)
 
 
@@ -579,8 +639,9 @@ def _sdf_narrowphase(
   mesh_polymapadr: wp.array(dtype=int),
   mesh_polymapnum: wp.array(dtype=int),
   mesh_polymap: wp.array(dtype=int),
-  volume_ids: wp.array(dtype=wp.uint64),
   oct_aabb: wp.array2d(dtype=wp.vec3),
+  oct_child: wp.array(dtype=vec8i),
+  oct_coeff: wp.array(dtype=vec8f),
   pair_dim: wp.array(dtype=int),
   pair_solref: wp.array2d(dtype=wp.vec2),
   pair_solreffriction: wp.array2d(dtype=wp.vec2),
@@ -732,11 +793,11 @@ def _sdf_narrowphase(
   rot1 = geom1.rot
 
   attr1, g1_plugin_id, volume_data1, mesh_data1 = get_sdf_params(
-    volume_ids, oct_aabb, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1]
+    oct_aabb, oct_child, oct_coeff, plugin, plugin_attr, type1, geom1.size, g1_plugin, geom_dataid[g1]
   )
 
   attr2, g2_plugin_id, volume_data2, mesh_data2 = get_sdf_params(
-    volume_ids, oct_aabb, plugin, plugin_attr, type2, geom2.size, g2_plugin, geom_dataid[g2]
+    oct_aabb, oct_child, oct_coeff, plugin, plugin_attr, type2, geom2.size, g2_plugin, geom_dataid[g2]
   )
 
   mesh_data1.nmeshface = nmeshface
@@ -857,8 +918,9 @@ def sdf_narrowphase(m: Model, d: Data):
       m.mesh_polymapadr,
       m.mesh_polymapnum,
       m.mesh_polymap,
-      m.volume_ids,
       m.oct_aabb,
+      m.oct_child,
+      m.oct_coeff,
       m.pair_dim,
       m.pair_solref,
       m.pair_solreffriction,

mujoco_warp/_src/io.py

@@ -96,7 +96,6 @@ def put_model(mjm: mujoco.MjModel) -> types.Model:
 
   plugin_id = np.array(plugin_id)
   plugin_attr = np.array(plugin_attr)
-  volume_ids, volumes, oct_aabb = _mujoco_octree_to_warp_volume(mjm)
 
   if mjm.nflex > 1:
     raise NotImplementedError("Only one flex is unsupported.")
@@ -637,9 +636,9 @@ def create_nmodel_batched_array(mjm_array, dtype, expand_dim=True):
     mesh_polymapadr=wp.array(mjm.mesh_polymapadr, dtype=int),
     mesh_polymapnum=wp.array(mjm.mesh_polymapnum, dtype=int),
     mesh_polymap=wp.array(mjm.mesh_polymap, dtype=int),
-    volume_ids=volume_ids,
-    volumes=volumes,
-    oct_aabb=oct_aabb,
+    oct_aabb=wp.array2d(mjm.oct_aabb, dtype=wp.vec3),
+    oct_child=wp.array(mjm.oct_child, dtype=types.vec8i),
+    oct_coeff=wp.array(mjm.oct_coeff, dtype=types.vec8f),
     nhfield=mjm.nhfield,
     nhfielddata=mjm.nhfielddata,
     hfield_adr=wp.array(mjm.hfield_adr, dtype=int),
@@ -858,85 +857,6 @@ def create_nmodel_batched_array(mjm_array, dtype, expand_dim=True):
   return m
 
 
-def _mujoco_octree_to_warp_volume(
-  mjm: mujoco.MjModel, resolution: int = 128
-) -> Tuple[wp.array, Tuple[wp.Volume, ...], wp.array]:
-  """Constructs volume data from MuJoCo octrees."""
-  volume_ids = [0] * len(mjm.mesh_octadr)
-  volumes = []
-  oct_aabbs = [None] * len(mjm.mesh_octadr)
-  for mesh_id in mjm.geom_dataid:
-    if mesh_id != -1 and mesh_id < len(mjm.mesh_octadr):
-      octadr = mjm.mesh_octadr[mesh_id]
-      octnum = mjm.mesh_octnum[mesh_id]
-      if octadr != -1:
-        oct_child = mjm.oct_child[octadr : (octadr + octnum), :]
-        oct_aabb = mjm.oct_aabb[octadr : (octadr + octnum), :]
-        oct_coeff = mjm.oct_coeff[octadr : (octadr + octnum), :]
-
-        root_aabb = oct_aabb[0]
-        center = root_aabb[:3]
-        half_size = root_aabb[3:]
-
-        original_mins = center - half_size
-        original_maxs = center + half_size
-
-        margin_factor = 0.02
-        extents = original_maxs - original_mins
-        margin = margin_factor * extents
-
-        mins = original_mins - margin
-        maxs = original_maxs + margin
-        expanded_extents = maxs - mins
-
-        voxel_size = expanded_extents.max() / resolution
-
-        nums = np.ceil(expanded_extents / voxel_size).astype(dtype=int)
-
-        actual_extents = nums * voxel_size
-        maxs = mins + actual_extents
-
-        sdf_values = np.zeros(tuple(nums), dtype=np.float32)
-
-        for x in range(nums[0]):
-          for y in range(nums[1]):
-            for z in range(nums[2]):
-              pos = mins + voxel_size * np.array([x, y, z])
-              within_bounds = np.all(pos >= original_mins) and np.all(pos <= original_maxs)
-              if within_bounds:
-                sdf_val = sample_octree_sdf(pos, oct_child, oct_aabb, oct_coeff)
-              else:
-                clamped_pos = np.clip(pos, original_mins, original_maxs)
-                sdf_val = sample_octree_sdf(clamped_pos, oct_child, oct_aabb, oct_coeff)
-              sdf_values[x, y, z] = sdf_val
-
-        device = wp.get_device()
-        if device.is_cuda:
-          volume = wp.Volume.load_from_numpy(sdf_values, mins, voxel_size, 1.0, device=device)
-          volume_ids[mesh_id] = volume.id
-          volumes.append(volume)
-        else:
-          volume_ids[mesh_id] = 0
-        oct_aabbs[mesh_id] = [center, half_size]
-
-  volume_ids_array = wp.array(data=volume_ids, dtype=wp.uint64)
-
-  processed_aabbs = []
-  for aabb in oct_aabbs:
-    if aabb is not None:
-      processed_aabbs.append(aabb)
-    else:
-      zero_center = np.zeros(3, dtype=np.float32)
-      zero_half_size = np.zeros(3, dtype=np.float32)
-      processed_aabbs.append([zero_center, zero_half_size])
-
-  aabb_array = np.array(processed_aabbs, dtype=np.float32)
-  oct_aabb_array = wp.array2d(data=aabb_array, dtype=wp.vec3)
-  volumes_tuple = tuple(volumes)
-
-  return volume_ids_array, volumes_tuple, oct_aabb_array
-
-
 def make_data(mjm: mujoco.MjModel, nworld: int = 1, nconmax: int = -1, njmax: int = -1) -> types.Data:
   """
   Creates a data object on device.

mujoco_warp/_src/io_test.py

@@ -130,9 +130,6 @@ def _check_annotation_compat(
     if isinstance(v, wp.types.array):
       continue
 
-    if v == wp.Volume:
-      continue
-
     if v in wp.types.vector_types:
       raise AssertionError(f"Vector types are not allowed. {info}")
 
@@ -392,30 +389,16 @@ def test_contact_sensor(self, contact_sensor):
     with self.assertRaises(NotImplementedError):
       mjwarp.put_model(mjm)
 
-  def test_volumes(self):
-    """Tests that mujoco_octree_to_warp_volume properly processes SDF volumes."""
-    if not wp.get_device().is_cuda:
-      self.skipTest("SDF volumes require CUDA device")
+  def test_sdf(self):
+    """Tests that an SDF can be loaded."""
     mjm, mjd, m, d = test_util.fixture(fname="collision_sdf/cow.xml", qpos0=True)
 
-    self.assertIsInstance(m.volume_ids, wp.array)
-    self.assertEqual(m.volume_ids.dtype, wp.uint64)
-    self.assertGreater(m.volume_ids.size, 0)
-
-    self.assertIsInstance(m.volumes, tuple)
-    if len(m.volumes) > 0:
-      for volume in m.volumes:
-        self.assertIsInstance(volume, wp.Volume)
-
     self.assertIsInstance(m.oct_aabb, wp.array)
     self.assertEqual(m.oct_aabb.dtype, wp.vec3)
     self.assertEqual(len(m.oct_aabb.shape), 2)
     if m.oct_aabb.size > 0:
       self.assertEqual(m.oct_aabb.shape[1], 2)
 
-    volume_ids_numpy = m.volume_ids.numpy()
-    self.assertEqual(len(m.volumes), np.unique(volume_ids_numpy).size)
-
 
 if __name__ == "__main__":
   wp.init()