diff --git a/mjx/mujoco/mjx/_src/ray.py b/mjx/mujoco/mjx/_src/ray.py index aaafc5a3b59..c22dad86b31 100644 --- a/mjx/mujoco/mjx/_src/ray.py +++ b/mjx/mujoco/mjx/_src/ray.py @@ -239,7 +239,7 @@ def ray( flg_static: bool = True, bodyexclude: Sequence[int] | int = -1, ) -> Tuple[jax.Array, jax.Array]: - """Returns the geom id and distance at which a ray intersects with a geom. + """Returns the distance and geom id at which a ray intersects with a geom. Args: m: MJX model @@ -289,7 +289,7 @@ def ray( dists, ids = dists + [dist], ids + [id_] if not ids: - return jp.array(-1), jp.array(-1.0) + return jp.array(-1.0), jp.array(-1) dists = jp.concatenate(dists) ids = jp.concatenate(ids) diff --git a/mjx/mujoco/mjx/_src/ray_test.py b/mjx/mujoco/mjx/_src/ray_test.py index ed0d6f19427..ad7fd07882e 100644 --- a/mjx/mujoco/mjx/_src/ray_test.py +++ b/mjx/mujoco/mjx/_src/ray_test.py @@ -48,6 +48,36 @@ def test_ray_nothing(self): _assert_eq(geomid, -1, 'geom_id') _assert_eq(dist, -1, 'dist') + def test_ray_no_geom_return_contract(self): + """ray()'s empty-filter early return keeps the (dist: float, id: int) contract.""" + m = test_util.load_test_file('ray.xml') + d = mujoco.MjData(m) + mujoco.mj_forward(m, d) + mx, dx = mjx.put_model(m), mjx.put_data(m, d) + + pnt, vec = jp.array([2, 1, 3.0]), jp.array([0.1, 0.2, -1.0]) + vec /= jp.linalg.norm(vec) + + # An all-zero geomgroup filters out every geom, so `ids` is empty and ray() + # takes its `if not ids:` early-return path (rather than the main path). + ray_fn = jax.jit(mjx.ray, static_argnums=(4,)) + dist_empty, geomid_empty = ray_fn(mx, dx, pnt, vec, (0, 0, 0, 0, 0, 0)) + + # The same pnt/vec hits the plane, exercising the main return path. + dist_hit, geomid_hit = jax.jit(mjx.ray)(mx, dx, pnt, vec) + + # Both output slots use the -1 sentinel when there is no intersection. + _assert_eq(dist_empty, -1, 'dist') + _assert_eq(geomid_empty, -1, 'geom_id') + + # Contract: distance is floating and geom id is integer, and the early + # return must agree with the main return on dtypes. Before the fix the early + # return is (int dist, float id) -- order and dtype swapped -- so this fails. + self.assertTrue(jp.issubdtype(dist_empty.dtype, jp.floating)) + self.assertTrue(jp.issubdtype(geomid_empty.dtype, jp.integer)) + self.assertEqual(dist_empty.dtype, dist_hit.dtype) + self.assertEqual(geomid_empty.dtype, geomid_hit.dtype) + def test_ray_plane(self): """Tests MJX ray<>plane matches MuJoCo.""" m = test_util.load_test_file('ray.xml')