Enable refit=True argument for render - #1559
Open
hartikainen wants to merge 2 commits into
Open
Conversation
`refit_bvh` mutates the scene BVH owned by the `RenderContext` and `render` reads it back, but the two exchange nothing that a graph executor can see, so nothing keeps a frame's refit paired with its own ray cast. In imperative Warp that is fine, since program order is the schedule, which is why `refit` defaults to `False` and no call site here changes. MJX's JAX/XLA bindings, however, lower the two to independent custom calls, and once several frames land in one executable XLA is free to hoist every refit ahead of every ray cast, so each frame ends up rendered against some other frame's BVH. `render(refit=True)` refits for itself, which makes a ray cast self-sufficient regardless of how the caller's executor orders things. MJX opts in, and nothing else has to.
`create_render_context` fits the BVH to the pose it is handed, and every existing render test renders that same pose, so none of them would notice `refit=True` quietly doing nothing. This adds a test that slides a box out from under the bounds fit at context creation, and then checks that `render(refit=True)` matches an explicit `refit_bvh` followed by `render`. It also guards itself by asserting that the stale BVH really does render a different image, since otherwise the test would pass whether or not the refit happens.
Collaborator
|
Hi @hartikainen - ouch, I can see this having resulted in a nasty bug. Hope it didn't bite you guys for very long. I have heard that Example here: Maybe there's a simple, generic way in the JAX wrapper code so that users can supply abstract tokens to indicate serial dependencies? cc @btaba |
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.
See google-deepmind/mujoco#3435 and google-deepmind/mujoco#3436.