Ray normal fixes (#960)

* Add ray with normal

* Ruff

* Propagate ray normal changes

* ruff

* Fix tuple type check

* Update ray to return -1 for non-hits

* Small fixes

* Ruff format

* Simplify assert for vec3

* Add comment and adjust docstring

Author: StafaH

mujoco_warp/_src/ray.py

@@ -393,7 +393,7 @@ def _ray_cylinder(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec:
 
 @wp.func
 def _ray_box(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, vec6, wp.vec3]:
-  """Returns the distance and normal at which a ray intersects with a box."""
+  """Returns distance, per side information, and normal at which a ray intersects with a box."""
   all = vec6(-1.0, -1.0, -1.0, -1.0, -1.0, -1.0)
 
   # bounding sphere test
@@ -812,9 +812,9 @@ def _ray(
   flg_static: bool,
   bodyexclude: wp.array(dtype=int),
   # Out:
-  dist_out: wp.array(dtype=float, ndim=2),
-  geomid_out: wp.array(dtype=int, ndim=2),
-  normal_out: wp.array(dtype=wp.vec3, ndim=2),
+  dist_out: wp.array2d(dtype=float),
+  geomid_out: wp.array2d(dtype=int),
+  normal_out: wp.array2d(dtype=wp.vec3),
 ):
   worldid, rayid, tid = wp.tid()
 

mujoco_warp/_src/ray_test.py

@@ -35,6 +35,7 @@ def _assert_eq(a, b, name):
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
 
+# TODO: Add tests comparing normal to engine implementation once available.
 class RayTest(absltest.TestCase):
   def test_ray_nothing(self):
     """Tests that ray returns -1 when nothing is hit."""
@@ -49,7 +50,7 @@ def test_ray_nothing(self):
     normal_np = normal.numpy()[0, 0]
     _assert_eq(geomid_np, -1, "geom_id")
     _assert_eq(dist_np, -1, "dist")
-    _assert_eq(normal_np, wp.vec3(0.0, 0.0, 0.0), "normal")
+    _assert_eq(normal_np, 0, "normal")
 
   def test_ray_plane(self):
     """Tests ray<>plane matches MuJoCo."""
@@ -292,7 +293,7 @@ def test_ray_bodyexclude(self):
     normal_np = normal.numpy()[0, 0]
     _assert_eq(geomid_np, -1, "geom_id")
     _assert_eq(dist_np, -1, "dist")
-    _assert_eq(normal_np, wp.vec3(0.0, 0.0, 0.0), "normal")
+    _assert_eq(normal_np, 0, "normal")
 
   def test_ray_invisible(self):
     """Tests ray doesn't hit transparent geoms."""
@@ -311,7 +312,7 @@ def test_ray_invisible(self):
     normal_np = normal.numpy()[0, 0]
     _assert_eq(geomid_np, -1, "geom_id")
     _assert_eq(dist_np, -1, "dist")
-    _assert_eq(normal_np, wp.vec3(0.0, 0.0, 0.0), "normal")
+    _assert_eq(normal_np, 0, "normal")
 
 
 if __name__ == "__main__":