Repository navigation
perf(mjwarp): skip the backend's blocking finite check under deferred… - #2598
Merged
TATP-233 merged 2 commits intoNov 7, 2026
Merged
Conversation
… validation MJWarp's step_tensor checked ctrl for NaN/Inf with a blocking device-to-host read on every call, duplicating TorchEnv's own control validation and making the host wait for the policy forward before it could submit physics. unisim-core 2.7.1 lets the caller own that check. - Pin unisim-core==2.7.1 in the default and ROCm profiles (both locks change only unisim-core). - Add EnvCfg.mjwarp_validate_tensor_finite (bool or None). None turns the backend check off when defer_finite_validation already validates on the device and keeps it otherwise; an explicit bool wins. env_backend_kwargs forwards it as mjwarp_validate_tensor_finite; other backends ignore it. On an RTX 5090, 50 step_tensor calls drop from 50 stream syncs and 50 DtoH copies to none under the default deferred validation. Refs Motphys#2587 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Make EnvCfg.mjwarp_validate_tensor_finite a bool defaulting to False and forward it unconditionally, so UniLab never inherits UniSim's blocking MJWarp default. TorchEnv validates the final control on every step (eagerly or through deferred finite validation), so the backend check is a duplicate in both modes; set the option to True only to debug a non-finite operand at the backend call. Refs Motphys#2587 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
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
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.
Summary
unisim-core==2.7.1in the default and ROCm profiles (pyproject.toml,pyproject.rocm.toml). In bothuv.lockanduv.rocm.lock, onlyunisim-corechanges.EnvCfg.mjwarp_validate_tensor_finite(bool orNone, defaultNone).env_backend_kwargspasses it to UniSim asmjwarp_validate_tensor_finite, and other backends ignore it.tests/base/backend/test_mjwarp_backend_options.py(13 tests).step_tensorcall, MJWarp checkedctrlfor NaN/Inf with a blocking device-to-host read. That duplicatesTorchEnv's own control check and makes the host wait for the policy forward pass before it can submit physics.mjwarp_validate_tensor_finite=None(the default), the MJWarp check is off whendefer_finite_validation=True, which is the default for training configs. It stays on whendefer_finite_validation=False. An explicittrue/falseoverrides this.Linked Work
mainValidation
Profiling:
torch.profilerover 50 MJWarpstep_tensorcalls (1024 envs, cartpole):defer_finite_validation=False(check on)defer_finite_validation=True(default, check off)These counts exclude the profiler's own end-of-trace
cudaDeviceSynchronize.Impact
mujoco/motrix. Onlymjwarpbehavior changes.Artifacts
Checklist