Skip to content

Render Cleanup - Ray#1081

Merged
thowell merged 5 commits into
google-deepmind:renderfrom
StafaH:render_cleanup_8
Jan 30, 2026
Merged

Render Cleanup - Ray#1081
thowell merged 5 commits into
google-deepmind:renderfrom
StafaH:render_cleanup_8

Conversation

@StafaH

@StafaH StafaH commented Jan 23, 2026

Copy link
Copy Markdown
Collaborator

This PR cleans up the following:

  • Allow for BVH accelerated ray and rays usage by passing in a RenderContext
  • Cleanup of imports to use absolute imports, matching rest of repo
  • Move BVH build into public API (fixes circular import issues, and makes it possible for downstream users to build BVH independently from the render context)
  • Add bvh ray tests

Comment thread mujoco_warp/__init__.py Outdated
from mujoco_warp._src.types import Data as Data
# isort: on

from mujoco_warp._src.bvh import build_scene_bvh as build_scene_bvh

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

@erikfrey @btaba what do we want to be part of the public api?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Why not expose bvh? Most of the other functions here seem to have an "equivalent in mujoco", so if we want to roughly keep that property, I propose bvh. Of course this is merely an opinion

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The bvh function is useful for downstream sensors that need to create or refit the BVH (e.g., raycast sensor in mjlab).

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I don't have a great intuition for which parts of the render code should be offered up in the public API, and at what granularity.

For example, if we are to expose both a refit_flex_bvh and a refit_scene_bvh then presumably there's some reason the user wants to refit just flexes? vs. a refit_bvh public function that just does it all?

Basically let's use making changes to the API in __init__.py as an invitation to discuss what the expected use cases are, maybe even with a few concrete user stories. Can discuss this here or in chat or in a doc, whatever you prefer @StafaH

In general we should be conservative about what we add to the API, the mistake of under-adding is much less painful to fix than the mistake of over-adding.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

Reduced to be just refit_bvh now

Comment thread mujoco_warp/_src/ray.py
Comment thread mujoco_warp/_src/render_context.py Outdated
from .types import GeomType
from .types import Model
from .types import ProjectionType
import mujoco_warp as mjw

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Let's use internal imports for internal code. So this could be:

from mujoco_warp._src import bvh

Comment thread mujoco_warp/_src/ray.py
Comment thread mujoco_warp/__init__.py Outdated
from mujoco_warp._src.types import Data as Data
# isort: on

from mujoco_warp._src.bvh import build_scene_bvh as build_scene_bvh

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I don't have a great intuition for which parts of the render code should be offered up in the public API, and at what granularity.

For example, if we are to expose both a refit_flex_bvh and a refit_scene_bvh then presumably there's some reason the user wants to refit just flexes? vs. a refit_bvh public function that just does it all?

Basically let's use making changes to the API in __init__.py as an invitation to discuss what the expected use cases are, maybe even with a few concrete user stories. Can discuss this here or in chat or in a doc, whatever you prefer @StafaH

In general we should be conservative about what we add to the API, the mistake of under-adding is much less painful to fix than the mistake of over-adding.

Comment thread mujoco_warp/_src/bvh.py
from mujoco_warp._src.types import Model
from mujoco_warp._src.warp_util import event_scope

if TYPE_CHECKING:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

is this check necessary? maybe add a brief comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

It stops a circular import between the two.

I think if we did a refactor to move render context into types and the init into the create method that already exists that would be cleaner and also resolve this

@thowell

thowell commented Jan 29, 2026

Copy link
Copy Markdown
Collaborator

left a question, otherwise lgtm. thanks @StafaH!

@thowell thowell merged commit 5f881a4 into google-deepmind:render Jan 30, 2026
1 check passed
erikfrey pushed a commit that referenced this pull request Feb 6, 2026
* 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>
@StafaH StafaH deleted the render_cleanup_8 branch February 8, 2026 03:47
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.

5 participants