Improve Startup Time - Newton Batched Model Builder#6430
Draft
StafaH wants to merge 6 commits into
Draft
Conversation
Replicating environments into the Newton model previously called ModelBuilder.add_builder once per environment and then rewrote entity labels in a second pass over the whole builder. Both steps are Python loops whose cost grows superlinearly with the environment count and dominate startup time for large scenes. Add BatchedModelBuilder and replicate_builder_mapping_batched, which consume the same clone-plan mapping inputs but append all worlds in one vectorized pass over cached per-source arrays, writing final per-environment labels directly so no rename pass is needed. Homogeneous and heterogeneous clone plans are supported; scenes with per-world builder hooks (e.g. MPM/deformable objects) fall back to the legacy path with a warning. The batched path is opt-in via NewtonCfg.use_batched_model_builder (or the explicit newton_physics_replicate_batched entry point); the legacy path remains the default and existing callers are unchanged. Also add compare_builder_states/compare_finalized_models for validating that two builders or finalized models are equivalent. For Isaac-Lift-KukaAllegro-Camera at 4096 envs, scene creation drops from 14.1s to 5.4s (replication itself from ~11s to ~1s).
Compare the legacy replicate+rename path against the batched path on identical inputs: builder state (all fields, exact), finalized models, site maps, world transforms, and Fabric body bindings. Covers a single world, many homogeneous worlds, heterogeneous variant mappings, multi-row worlds, worlds without rows, env-root and injected source sites, destination builders with pre-existing global entities, and an end-to-end Isaac-Lift-KukaAllegro-Camera scene with the newton_mjwarp preset. Also verifies the vectorized float32 transform composition is bitwise-identical to Warp's builtins.
benchmark_model_builder.py times build, finalize, and total time for the legacy and batched Newton replication paths across configurable tasks and world counts, printing a comparison table with per-row speedups and writing JSON/CSV. Each task runs in its own subprocess (Isaac Sim can launch once per process); --verify cross-checks that both paths produce equivalent builders.
replicate_builder_mapping built the per-label site map via setdefault, whose default argument allocates a num_worlds list on every call - once per (world, site label) - making startup quadratic in the environment count when sites are registered. Precompute the map for all labels once. Also pass the configured physics device to wp.ScopedCapture so CUDA graph capture does not select the default GPU on multi-GPU machines.
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
1 participant
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.
Description
Checklist
pre-commitchecks with./isaaclab.sh --formatconfig/extension.tomlfileCONTRIBUTORS.mdor my name already exists there