Fix rendering race condition - #3436
Open
hartikainen wants to merge 3 commits into
Open
Conversation
hartikainen
force-pushed
the
fix-rendering-race-condition
branch
from
July 27, 2026 14:46
3731ccf to
6f8a98c
Compare
Contributor
|
@hartikainen thanks for reporting this issue and submitting a pr to fix it! it would be great if we could resolve this without needing to change mujoco warp since my understanding is that the reported issue is related to the way jax optimizes jitted code and not the source code itself. |
Contributor
Author
Correct. I think we need to somehow tie CC @tkelestemur @mhasek. |
`mjx.refit_bvh` mutates the scene BVH owned by the render context and `mjx.render` reads it back, but the two lower to custom calls that exchange nothing through the JAX graph, so XLA is free to order them however it likes. Rendering one frame per executable hides this, since JAX dispatches executables in order. Rendering several frames from one executable does not, and XLA hoists every refit ahead of every ray cast, so each frame ends up rendered against some other frame's BVH.
`mujoco_warp`'s `render` gains a `refit` argument that makes a ray cast self-sufficient, and the generated shim passes `refit=True` so MJX no longer depends on the scheduler pairing the two calls. The argument defaults to `False`, which leaves the imperative Warp pipelines that already call `refit_bvh` themselves unchanged.
Regenerating `mjx/mujoco/mjx/warp/render.py` also picks up the ten fields that `refit_{bvh,scene_bvh,flex_bvh}` read and the shim was missing (`nflex`, `nflexelem`, `flex_vertnum`, `flex_dim`, `flex_elem`, `flex_elem{adr,dataadr,num}`, `flex_shell`, and `flex_shelldataadr`). Without them a model with a `flexcomp` hands Warp a null `flex_vertnum` and segfaults. I left the generated import block as it is checked in, since the local `isort` and `pyink` reshuffle it in ways unrelated to this change.
The vendored `mujoco_warp` hunk is carried here until the same change lands upstream, after which a re-vendor supersedes it.
mjx.refit_bvh writes the scene BVH owned by the render context and mjx.render reads it back, but the two lower to custom calls that exchange nothing through the JAX graph, so nothing keeps a frame's refit paired with its own ray cast. Rendering several frames from one executable makes XLA hoist every refit ahead of every ray cast, and those frames come back rendered against another frame's bounding volumes. Give the shim generator a --mjwarp_prelude_function flag that merges a second mjwarp function's field usage into the shim signature and calls it first, then use it to fuse bvh.refit_bvh into the render shim. The refit and the ray cast become a single FFI call that nothing can interpose between, and mujoco_warp itself is untouched. The merged field usage also brings in the ten flex fields the refit reads and the render shim did not previously carry. mjx.refit_bvh stays public for pipelines that use the BVH without rendering, but it is now redundant before a render, so drop it from the docs and from the two in-repo callers. Add two regression tests: one renders five poses one per executable and again from a single executable, the other deliberately leaves the BVH fit to a different pose before rendering. Fixes google-deepmind#3435.
hartikainen
force-pushed
the
fix-rendering-race-condition
branch
from
July 28, 2026 09:56
6f8a98c to
5c1fee6
Compare
hartikainen
commented
Jul 28, 2026
hartikainen
left a comment
hartikainen
left a comment
Contributor
Author
There was a problem hiding this comment.
@thowell I drafted a fix for this with Claude. It required some weird prelude stuff for the warp shim. I can clean those up and make sure they're sound if you think the overall approach of fusing the refit_bvh and render together is reasonable.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #3435.