Add ray with normal by StafaH · Pull Request #940 · google-deepmind/mujoco_warp

StafaH · 2025-12-21T10:09:40Z

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.

thowell · 2025-12-23T17:21:13Z

thanks @StafaH!

erikfrey

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

erikfrey · 2025-12-27T19:28:26Z

        x = all[i]

-  return x
+        normal_local = wp.vec3(0.0, 0.0, 0.0)


how about:

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

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

erikfrey · 2025-12-27T19:30:33Z

    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)


nit: wp.vec3() already inits to zeros

(comment applies here and elsewhere)

thowell · 2025-12-28T10:30:52Z

+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."""


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

thowell · 2025-12-28T10:31:53Z

 @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."""


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

thowell · 2025-12-28T10:32:22Z

 @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."""


update docstring

StafaH · 2026-01-01T21:18:10Z

@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.

thowell

added a few more comments. otherwise lgtm.

great work and thank you for this contribution @StafaH!

thowell · 2026-01-02T13:40:52Z

+    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")


wp.vec3 -> 0.0

* Add ray with normal (#940) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * 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 * Fix bugs in ray hfield impl (#967)

* Add ray with normal (#940) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * 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 * Fix bugs in ray hfield impl (#967) * clamp pair_friction with minmu (#969) * add --kernel_cache_dir pytest flag for warp kernel cache (#968) * fix _NCONMAX info (#965) * reuse xnorm (#959) * Expose rays through the public API. (#975) --------- Co-authored-by: Taylor Howell <taylorhowell@google.com> Co-authored-by: Kevin Zakka <kevinzakka@users.noreply.github.com>

* Render (#836) * Integrate madrona_warp into mjwarp * Fix bug for <8 views * Move mesh construction to SAH, tile pixels * Change to registry to pass jax io test * Remove dist from bvh_query_ray * Add render context * Update testspeed * Add render context registry wrapper * Fix circular import * fix ray box and ambient lighting * Cleanup rc options in kernel * Small fix to static nlight * ray capsule fix * Fix testspeed merge artifact * Cleanup before PR into mjwarp * Remove global fovy, use cam 0 values * Cleanup bvh/render/ray * Fix capsule calc comment * Cleanup docstrings * Cleanup testspeed for render * Fix bug when no texid is set * Render hetero (#873) * Integrate madrona_warp into mjwarp * Fix bug for <8 views * Move mesh construction to SAH, tile pixels * Change to registry to pass jax io test * Remove dist from bvh_query_ray * Add render context * Update testspeed * Add render context registry wrapper * Fix circular import * fix ray box and ambient lighting * Cleanup rc options in kernel * Small fix to static nlight * ray capsule fix * Fix testspeed merge artifact * Cleanup before PR into mjwarp * Remove global fovy, use cam 0 values * Cleanup bvh/render/ray * Fix capsule calc comment * Cleanup docstrings * Cleanup testspeed for render * Fix bug when no texid is set * hetero camera support * Add heightfield support * Add flex mesh building * Add correct flex 2d mesh building * Fixes for downstream use * Rename output buffers to match mujoco style * Add flex support, improve render script * Add selective camera rendering support via cam_active parameter * Fix render output reading, fix bg color * Add more optimzed hfield mesh building * Add visual franka benchmark, fix merge artifacts * Add 3d flex rendering support * Add ellipsoid and cylinder primitives * Cleanup flex 2d and 3d rendering * Cleanup of bvh code * Fix 2d flex side triangle update * More cleanup for PR * More cleanup of naming, paranthesis, casts * Add correct frustrum calculation for rays * Ruff check * Fix naming, add todo * Update kernel arguments using analyzer * Remove comment * Remove visual assets * Add back camera * Remove stray 1, dim --------- Co-authored-by: Kevin Zakka <kevinarmandzakka@gmail.com> * Render merge ray normal (#979) * Add ray with normal (#940) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * 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 * Fix bugs in ray hfield impl (#967) * Merge main into render branch (#980) * Add ray with normal (#940) * Add ray with normal * Ruff * Propagate ray normal changes * ruff * Fix tuple type check * Update ray to return -1 for non-hits * 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 * Fix bugs in ray hfield impl (#967) * clamp pair_friction with minmu (#969) * add --kernel_cache_dir pytest flag for warp kernel cache (#968) * fix _NCONMAX info (#965) * reuse xnorm (#959) * Expose rays through the public API. (#975) --------- Co-authored-by: Taylor Howell <taylorhowell@google.com> Co-authored-by: Kevin Zakka <kevinzakka@users.noreply.github.com> * Fix render aspect ratio bug (#1002) * Update to new cam proj field * Fix frustrum aspect ratio * Render tests and cleanup (#1021) * ruff * Move render script to contrib * Add render tests * Move create render context into io.py * Cleanup bvh test, ruff * cleanup render_context_test * More cleanup * Render Cleanup 2 (#1037) * Update context to read from new mujoco fields * Add anyhit ray mesh query for fast shadows * Cleanup tests, msgs, comments * Render Cleanup 3 (#1063) * Fix naming, move mesh bvh build to bvh module * Add batched cam intrinsic for randomization * Fix indexing in render context * ruff * Render Cleanup 4 (#1064) * Fix failing tests with new batched cam fields * Move hfield bvh build to bvh module * Fix docstring * Render Cleanup 5 (#1075) * Fix flex rendering for multiple flex * Add flex bvh test * Render Cleanup - Textures (#1080) * Add cuda texture sampling * Fix textures and mesh sample, cleanup * ruff * Clean up texture processing kernel * Fix flipped texture uv * Remove nested texture kernel * Render Cleanup - Ray (#1081) * Add bvh accelerated ray * Move ray_bvh into ray * Remove ray_bvh_test and merge with ray_test * Use absolute import * Cleanup public API for render/bvh * Render Cleanup - types and io (#1103) * Move RenderContext to types * Cleanup flex fields and refit kernel * Use array in render context type * Fix all render related tests after refactor * Move packing to render util * Cleanup render file * Cleanup based on pr comments * Fix error with rc init * Fix flex render error (#1108) * Update _render_megakernel decorator (#1109) * Fix error for downstream mjlab * Fix ray none check * Change nested kernel to wp.kernel * Remove TODO * Revert uv.lock * Post merge cleanup and fixes * Add guards for bleeding edge features * Fix guards for mjtProjection * Fix failing test, post merge artifact * Fix failing tests for ray/render * Add more guards for mujoco < 3.4.1 * Ruff * Disable flex cpu test, fix leading dim * Add newline --------- Co-authored-by: Kevin Zakka <kevinarmandzakka@gmail.com> Co-authored-by: Taylor Howell <taylorhowell@google.com> Co-authored-by: Kevin Zakka <kevinzakka@users.noreply.github.com>

Add ray with normal

7299911

StafaH requested a review from thowell December 21, 2025 10:09

StafaH added 4 commits December 21, 2025 15:12

Ruff

ff4c354

Propagate ray normal changes

471cd08

ruff

9f4b64d

Fix tuple type check

86787c9

erikfrey reviewed Dec 27, 2025

View reviewed changes

thowell reviewed Dec 28, 2025

View reviewed changes

Update ray to return -1 for non-hits

1e67cd1

StafaH requested review from erikfrey and thowell January 1, 2026 21:18

thowell approved these changes Jan 2, 2026

View reviewed changes

thowell marked this pull request as ready for review January 2, 2026 18:55

thowell merged commit c915c40 into google-deepmind:main Jan 2, 2026
13 of 16 checks passed

Conversation

StafaH commented Dec 21, 2025

Uh oh!

thowell commented Dec 23, 2025

Uh oh!

erikfrey left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

StafaH commented Jan 1, 2026

Uh oh!

thowell left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants