Skip to content

Conversation

@StafaH
Copy link
Collaborator

@StafaH StafaH commented Dec 21, 2025

Modify ray functions to return normal information as well.

Follows the C MuJoCo implementation: https://github.com/google-deepmind/mujoco/blob/main/src/engine/engine_ray.c

Currently waiting on the C MuJoCo mju_rayNormal to become enabled, which will allow ray_test to compare normals with the engine implementation.

@StafaH StafaH requested a review from thowell December 21, 2025 10:09
@thowell
Copy link
Collaborator

thowell commented Dec 23, 2025

thanks @StafaH!

Copy link
Collaborator

@erikfrey erikfrey left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looking good! Just a few drive-by christmas nits.

x = all[i]

return x
normal_local = wp.vec3(0.0, 0.0, 0.0)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

how about:

normal_local = wp.vec3(float(i == 1) - float(i == 0), float(i == 3) - float(i == 2), 0.0)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is a very nifty one-liner! Beautiful. Replaced.

or (planar_01 < 0.0 and planar_11 < 0.0 and planar_21 < 0.0)
):
return float(wp.inf)
return float(wp.inf), wp.vec3(0.0, 0.0, 0.0)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: wp.vec3() already inits to zeros

(comment applies here and elsewhere)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

def _ray_triangle(
v0: wp.vec3, v1: wp.vec3, v2: wp.vec3, pnt: wp.vec3, vec: wp.vec3, b0: wp.vec3, b1: wp.vec3
) -> Tuple[float, wp.vec3]:
"""Returns the distance at which a ray intersects with a triangle."""
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's update the docstring to mention that the normal (?) is also returned

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

@wp.func
def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> float:
def _ray_plane(pos: wp.vec3, mat: wp.mat33, size: wp.vec3, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
"""Returns the distance at which a ray intersects with a plane."""
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

let's update the docstring to mention that the normal is returned

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

@wp.func
def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> float:
def _ray_sphere(pos: wp.vec3, dist_sqr: float, pnt: wp.vec3, vec: wp.vec3) -> Tuple[float, wp.vec3]:
"""Returns the distance at which a ray intersects with a sphere."""
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

update docstring

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed

@StafaH
Copy link
Collaborator Author

StafaH commented Jan 1, 2026

@thowell @erikfrey Fixed up the comments. Also moved to returning -1 for non-hits to match MuJoCo implementation. Just waiting now on C MuJoCo to enable the ray + normal and we can add tests to make sure the normals match the engine impl.

@StafaH StafaH requested review from erikfrey and thowell January 1, 2026 21:18
Copy link
Collaborator

@thowell thowell left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

added a few more comments. otherwise lgtm.

great work and thank you for this contribution @StafaH!

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")
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

wp.vec3 -> 0.0

@thowell thowell marked this pull request as ready for review January 2, 2026 18:55
@thowell thowell merged commit c915c40 into google-deepmind:main Jan 2, 2026
13 of 16 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants