Integrate articulation ordering caches and selectors - #6784
Conversation
The articulation element kind was a stringly-typed Literal["joint", "body"] threaded through the ordering resolvers, with all kind-dependent data scattered as inline ternaries and per-kind module dicts. Introduce a module-private frozen _ArticulationElementKind dataclass that carries the per-kind data as fields: the backend-names attribute, the robot-schema relationship name and name-override attribute spellings, the config-field label, and the spelling-match flag. Two interned singletons (_JOINT_KIND, _BODY_KIND) and a label registry replace the scattered ternaries and the _ROBOT_SCHEMA_RELATIONSHIP_NAMES and _ROBOT_SCHEMA_NAME_OVERRIDE_ATTRS dicts. _validate_ordering_kind becomes the single boundary conversion point, accepting the public "joint"/"body" string (or an already-resolved kind) and returning the singleton. The public entry points keep accepting the string, so the external contract is unchanged. The per-articulation convention cache stays keyed by the string label to preserve the cache-key type.
Translate between public and backend order on the OVPhysX read and write paths: fused user-and-backend writes with persistent target staging separated from transient applied-torque scratch, forward kinematics aware body state refresh, and a fused ordered joint acceleration kernel. Part 6 of the articulation reordering series.
The ordering core rework single-homed the joint and body ordering maps on ArticulationData and removed the mirrored `_has_*` flags and `_*_user_to_backend` maps, the `_cache_ordering_maps` machinery, the `ArticulationNameMap` identity/name fields, and the dtype-suffixed write-kernel overloads. Repoint the OVPhysX articulation and data classes to the new contracts: - Read ordering presence via `data.joint_ordering` / `data.body_ordering` is-not-None checks; conditional sites dereference the map directly, while always-launch sites use the new `_*_user_to_backend_map()` helpers so publish, wrench, and telemetry gathers keep a valid map. - Configure backend staging through `_ordering_configure_backend_staging` and, now that `_install_ordering_flags` is gone, recover the previous per-axis ordering from the persistent backend staging buffers so a cleared ordering still releases its buffers on rebind. - Replace the dtype-suffixed overload launches with the `write_*_user_to_backend_*` launch wrappers, passing the dtype the old suffix implied and dropping the explicit launch dim. Adapt the OVPhysX asset tests to assert through the single-homed ordering maps, since a non-None map now always denotes a real permutation.
Add the shared interface test suite that exercises ordering allocation, write parity, partial writes, and property reordering uniformly across the PhysX, OVPhysX, and Newton backends through their mock views. Lands after all three backend parts because present-but-unwired backends error rather than skip. Part 7 of the articulation reordering series.
The ordering core now keeps a single source of truth for public ordering. Identity orderings resolve to None rather than an identity map, so a non-None map always denotes a real permutation. The mirror flags (_has_joint_ordering / _has_body_ordering on both the asset and data), the _cache_ordering_maps / _ordering_maps_cached guard, and the data _install_ordering_flags helper were all deleted; backends now stage through _ordering_configure_backend_staging(). Adapt the mocked cross-backend interface tests to match: - Re-stage installed orderings via _ordering_configure_backend_staging instead of the removed _cache_ordering_maps guard dance. - Delete the reinvocation-guard test; with one home there is no mirror state left to protect. - Assert ordering presence as (ordering is not None) against the single source instead of the removed mirror flags, keeping one dedicated assert that the asset property delegates to data. - Drop the now-redundant not-is_identity assert halves, since a non-None map already guarantees a permutation.
Keep the backend import-semantics regression with the shared interface utility tests introduced in P5.
Route event terms through the asset's public index mapping, fix consumers that mixed public-order indices with backend-order view arrays, apply the factory inertia offsets through the asset API, and document the feature: the sim-to-sim transfer guide, the API reference entries, and the changelog fragments for the touched task packages. Part 8 of the articulation reordering series.
A 10-seed paired benchmark measured the feature at 1.06-1.55% lower native training FPS than develop, attributed to write_data_to_sim launching the joint-target gather unconditionally with an identity map. This reverses the unification from 470b4b7 on the strength of that measurement: both actuator branches now branch once in Python and take develop's exact direct-assign path when no joint ordering is configured, and the external-wrench path selects the reorder-free shared kernel when body ordering is identity instead of threading a flag through the ordered kernel every step. Also fuse the post-step publish hook from four reorder launches into two: one kernel gathers joint position and velocity together, one gathers body pose and velocity together.
The ordering rework removed six public kernels without a prior deprecation release: write_joint_state_data and write_joint_vel_data on PhysX, and the four write_joint_state_data and write_joint_vel_data index/mask variants on Newton. No ordering-aware replacement is signature compatible, so each is restored verbatim from develop as a deprecated shim for one release, with the replacement APIs named in the docstrings and changelog. Raw warp kernels cannot emit a runtime warning without breaking wp.launch, so the deprecation is documented rather than raised.
Resolution errors told users to configure env.scene.robot, which is wrong for direct-workflow tasks. Point at the ArticulationCfg fields instead, which are layout independent.
Drop the two map micro-tests that duplicate the gather-branch coverage of the permutation tests in the same file and the name-map construction coverage in the core ordering suite. The four remaining map tests are the only coverage of the public map_body_ids_to_backend and map_joint_ids_to_backend application paths and stay.
Replace the 108-case partial-write cartesian with a deterministic 12-row all-pairs covering array, verified to cover every parameter pair and seeded on all backend, ordering, and selector triples so each backend still drives both writer kinds under both orderings. Drop the identity value from the shadow-allocation axes, covered by the per-backend identity-matches-none tests, and remove the internal buffer-shape assertions whose behavioral observables are covered by the shared-backend-read and allocation tests. 233 collected cases become 131 with no loss of distinct-branch coverage.
Application coverage of map_body_ids_to_backend and map_joint_ids_to_backend now lives beside the name-map construction tests in the core ordering suite.
In the ordered Lab-actuator branch of write_data_to_sim, the fused target reorder passed the effort buffer as both the effort output (index 0) and the unused joint-act output (index 3). Because write_joint_act is off, index 3 is never written, but the alias creates a false output dependency on the effort buffer under CUDA graph capture. Point index 3 at the dedicated _sim_bind_joint_act buffer, matching the newton-actuator branch.
sync_torque_telemetry mixed the user-space loop index j with the backend-space backend_j. Rename the user-space index to user_j so each index only subscripts arrays of its own order: user-order targets, gains, limits, and outputs at [i, user_j]; backend-order live state and sim-bound effort buffers at [i, backend_j]. The kernel signature is unchanged, so call sites are unaffected.
Document that register_post_step_callback's hook fires exactly once per step() call, reflecting state after all decimation iterations and their solver substeps complete -- not once per substep or per decimation iteration.
The remaining public-order presence re-derivations in write_data_to_sim tested the raw ordering map (self.data.joint_ordering is None). Replace them with the has_joint_ordering property so presence checks read uniformly across the ordering paths. Genuine map-object uses and articulation-level ordering reads are left untouched. No behavior change.
Record stable PhysX ordering reads once and replay their Warp commands. Avoid nested recording during Newton CUDA graph capture, and discard cached commands when articulations reset or ordering buffers are rebound.
Cache stable PhysX and Newton articulation read kernels so repeated property refreshes avoid Python and Warp launch setup. Preserve eager execution for dynamic scalar inputs and invalidate commands only when their captured buffers can change.
Warp returns no recorded command for zero-sized launch dimensions. Treat that result as a completed no-op so fixed and zero-DOF articulations can initialize with read caching enabled.
PhysX-family joint states follow authored USD direction while dynamics tensors use the backend topology direction. Cache that difference at startup and normalize dynamics reads into the public joint basis.
marcodiiga
left a comment
marcodiiga
left a comment
There was a problem hiding this comment.
Thanks for slimming this down and addressing the earlier findings. I re-reviewed the current head. The earlier actuator, SceneEntityCfg, and WrenchComposer issues look addressed, but I found four blocking correctness/CI issues: simulator-bound int64 selectors, cached CUDA launches after buffer rebinding, a Sphinx import failure, and instance-proxy traversal for reversed joints. Details are inline.
kellyguo11
left a comment
kellyguo11
left a comment
There was a problem hiding this comment.
have we added some documentation for the reordering? when it's useful, how to use it...
There was a problem hiding this comment.
Agent reviews: Requesting changes for three additional CI blockers called out inline. The four existing unresolved threads covering simulator-bound int64 selectors, recorded CUDA launches after buffer rebinding, Sphinx annotation evaluation, and instance-proxy traversal remain blocking as well. Pre-commit and the pure selector/cache utility tests pass, but the current integration matrix is not merge-ready.
| self._fk_timestamp = 0.0 | ||
| self._cached_read_launches: dict[object, wp.Launch] = {} | ||
| self._read_launch_cache = _WarpLaunchCache(device) | ||
| self._joint_dof_signs = wp.ones(root_view.max_dofs, dtype=wp.int32, device=device) |
There was a problem hiding this comment.
Could we keep the shared PhysX view contract in sync with this new property, or derive the DoF count from an existing contract such as shared_metatype.dof_count? MockArticulationViewWarp does not expose max_dofs, so ArticulationData construction now fails before the shared PhysX interface tests can run. Across the current core shards this produces 1,188 failures with AttributeError: 'MockArticulationViewWarp' object has no attribute 'max_dofs'. Please update the mock and cover construction through the shared interface suite.
There was a problem hiding this comment.
Fixed in 19e72de. MockArticulationViewWarp now exposes max_dofs from its existing DOF count, and its initialization test asserts max_dofs alongside shared_metatype.dof_count. The focused PhysX/OVPhysX kernel and mock-interface group passes 61 tests.
| root_transforms[..., 3:7] /= np.linalg.norm(root_transforms[..., 3:7], axis=-1, keepdims=True) | ||
| mock_view.set_mock_root_transforms(wp.array(root_transforms, dtype=wp.transformf, device=device)) | ||
| mock_view.set_mock_root_velocities( | ||
| mock_view._root_transforms.assign(wp.array(root_transforms, dtype=wp.transformf, device=device)) |
There was a problem hiding this comment.
Could we preserve the mock-view setter interface while keeping the backing arrays stable? This refresh helper now reaches directly into _root_transforms, _root_velocities, and _attributes, but the supported recording view exposes set_mock_* methods instead. Consequently both Newton benchmark runtime regressions fail, beginning with test_newton_rigid_object_refresh_regenerates_kinematics_and_mass_properties at this line. If pointer stability is the goal, the setters can assign into stable buffers so callers do not need to depend on private storage.
There was a problem hiding this comment.
Fixed in 19e72de. The Newton refresh helpers use the set_mock_* interface; articulation setters assign into lazily allocated stable buffers, and the collection mock exposes the same stable setter surface. The two runtime regressions and the 41-test mock-view suite pass.
| import torch | ||
| import warp as wp | ||
|
|
||
| from .index_kernel import IndexKernelDispatcher |
There was a problem hiding this comment.
Could we preserve the standalone loading path used by test_ray_caster_kernels.py? That test loads this file through spec_from_file_location, so this new relative import runs without a package context and raises ImportError: attempted relative import with no known parent package during collection. Please either make the dispatcher import standalone-safe or update the test to import through the package without reintroducing its heavyweight runtime dependencies.
There was a problem hiding this comment.
Fixed in 19e72de with the absolute isaaclab.utils.warp.index_kernel import. The standalone test remains unchanged and all 12 ray-caster kernel tests pass on CPU.
Materialize simulator-compatible environment indices in the existing joint write kernels and reuse per-articulation scratch storage.
Restore the baseline before measuring the mask alternative with the same benchmark.
Materialize environment masks in the existing joint write kernels and reuse per-articulation mask storage at simulator boundaries.
Replace the mask prototype with the faster fused int32 scratch path measured by the articulation micro-benchmark.
Mirror the Jacobian, mass matrix, and gravity-force bindings created by real articulations so the data micro-benchmark completes.
Remove the obsolete full-data flag from the fused joint-state scratch launch so int64 selectors reach the simulator boundary correctly.
|
Yes. This PR includes the sim-to-sim policy transfer guide. It covers when joint/body reordering is useful, configuration and CLI overrides, ordering presets, public/backend mapping semantics, and raw backend access. |
Preserve selector dtype dispatch while normalizing simulator-bound indices. Keep cached reads and benchmark mocks stable across PhysX, OVPhysX, and Newton, and handle reversed joints in instanceable OVPhysX assets.
|
@marcodiiga @kellyguo11 All outstanding inline findings are addressed on head 19e72de, with a response on each thread. Focused verification is 131 passed and 1 skipped, the mocked Sphinx import passes, and the full pre-commit hook set is clean. Ready for re-review. |
marcodiiga
left a comment
marcodiiga
left a comment
There was a problem hiding this comment.
Re-reviewed exact head 19e72debaaf66573f1a0040ad5ffe340300dbd1f. The nine issues from my earlier two reviews are addressed. I found two remaining regressions, detailed inline: the fused simulator-index path mishandles empty joint selections, and one new PhysX dispatcher annotation still evaluates under Sphinx mocks.
Return after validating empty environment or joint selections so stale simulator index scratch buffers cannot trigger backend writes.
Cover list, Torch, Warp, and mixed-width selector modes. Run mocked Newton and PhysX asset suites kitless so startup work stays outside the measured asset path.
Emit int32 simulator selectors from existing indexed work kernels and reuse pinned CPU prefixes at PhysX and OVPhysX boundaries. Preserve int32 and int64 selectors through the dispatched work path without per-call conversion allocations.
marcodiiga
left a comment
marcodiiga
left a comment
There was a problem hiding this comment.
Re-reviewed exact head 00c8c25cb102309f550081fb06ca352cd446055d. The two issues from my previous review are cleared: empty joint selections are now no-ops in both backends, and the supported warnings-as-errors documentation build completes with zero warnings at this SHA. I found five follow-up issues in the latest scratch, benchmark, and documentation delta and left them inline. Approving this revision; the broader CI matrix is still in progress.
| if count not in self._sim_view_ids_views: | ||
| self._sim_view_ids_views[count] = wp.array( | ||
| ptr=self._sim_view_ids.ptr, | ||
| shape=(count,), |
There was a problem hiding this comment.
Could we bounds-check count against the backing allocation before creating this view? _sim_view_ids is sized to num_instances * num_bodies, while count is derived from selector lengths and the public APIs accept repeated selectors. On a 2x3 collection, env_ids=[0, 0, 0, 1, 1] with all bodies requests a 15-element view over a six-element allocation; the CPU repro terminated with a core dump and CUDA reported an invalid argument. The previous path allocated exactly count. Please fall back to a correctly sized allocation or reject oversized selectors, and apply the same guard to the mirrored prefix-view helpers.
There was a problem hiding this comment.
Fixed in a83b390. Selector cardinality is now checked under the existing shape-validation gate before any prefix view or launch. This rejects selector lengths beyond the asset dimensions (environment IDs are required to be unique) and applies the environment/body bounds through the shared asset bases. The existing PhysX/OVPhysX collection scratch regression now covers the oversized case.
| if device != "cpu" or self.device == "cpu": | ||
| return view_ids | ||
| cpu_view_ids = self._cpu_view_ids_view(count) | ||
| wp.copy(cpu_view_ids, view_ids) |
There was a problem hiding this comment.
Could we establish completion before returning this pinned host view? wp.copy from CUDA into pinned host memory is asynchronous on Warp's current stream, and this result is handed directly to a CPU simulator setter. With the destination prefilled, I observed the previous values immediately after wp.copy and the expected indices only after synchronization. In this path, the blocking data reshape happens before this copy, so it does not order the subsequent host read. Please add an explicit stream/event handoff or use a synchronous ownership boundary; the OVPhysX mirror has the same ordering.
There was a problem hiding this comment.
Fixed in a83b390. Both PhysX and OVPhysX now synchronize Warp's current device stream immediately after the CUDA-to-pinned-host index copy and before returning the CPU view. The existing mirrored collection test verifies copy-before-sync ordering. A matched 1,000-iteration GPU full-grid comparison at 4,096 instances measured a localized ~1–3% aggregate cost across masses, CoMs, and inertias.
| Asset Input Modes | ||
| ~~~~~~~~~~~~~~~~~ | ||
|
|
||
| Methods indexed only by ``env_ids`` retain two modes: |
There was a problem hiding this comment.
Could we update this mode contract and validate unknown --mode values? The environment-only generators now emit five modes and no torch_tensor; the item generators emit nine modes, and several vary both env_ids and the item selector rather than holding environment IDs at Torch int32. As written, the documented --mode torch_tensor matches no workloads and the CLI then fails during finalization with KeyError: 'mean'. The sample [TORCH_TENSOR] output is stale as well. Please align the names, grid, and docs, and reject unmatched modes so commands cannot benchmark nothing.
There was a problem hiding this comment.
Fixed in a83b390. The documentation now describes the five environment-selector modes and nine paired item-selector modes, and the sample output uses a current mode name. The CLI derives exact choices from the selected backend/component definitions, so an unmatched mode exits through argparse instead of reaching empty finalization; this case was folded into the existing invalid-argument test.
| - Sensor benchmarks | ||
| * - PhysX | ||
| - Launch Isaac Sim and use mocked PhysX views | ||
| - Run kitless with mocked PhysX views |
There was a problem hiding this comment.
Could we qualify this requirement? The prerequisites table now says PhysX asset benchmarks run kitless with mocked PhysX views, and the asset runtime no longer launches the app. Isaac Sim is still required for the PhysX sensor benchmarks, but not for the asset benchmarks described above.
There was a problem hiding this comment.
Fixed in a83b390. The troubleshooting text now states that PhysX sensor benchmarks require Isaac Sim; the kitless PhysX asset benchmark requirement remains as documented in the prerequisites table.
| "omni.physics": omni_physics, | ||
| "omni.physics.tensors": omni_physics.tensors, | ||
| } | ||
| with patch.dict(sys.modules, modules): |
There was a problem hiding this comment.
Could we avoid restoring the entire sys.modules mapping after these imports? patch.dict(sys.modules, modules) removes every module first imported inside the block when it exits. After _load_runtime_symbols(), isaaclab_physx.physics and the imported asset modules are no longer registered; re-importing the articulation module in the same kitless process then fails on the real carb dependency, while the globals retained here still refer to the first class objects. Limiting cleanup to the injected mock keys would avoid failed later imports and duplicate class identities.
There was a problem hiding this comment.
Fixed in a83b390. The runtime now restores only the four injected mock-module keys and leaves modules imported during setup registered. The existing runtime test now removes the articulation module, reloads the symbols, and verifies that the registered module retains the same Articulation class identity.
Bound selector scratch use to asset dimensions, synchronize pinned-host index copies before simulator writes, and keep kitless PhysX imports registered. Align benchmark CLI validation and documentation with the supported mode grid.
Initialize selector scratch state in interface fixtures so mocked assets match production initialization. Pass resolved body selectors to the inertia kernel to preserve dtype dispatch and partial selections.
Initialize view-index scratch state in the mocked collection factory so interface tests follow production buffer setup.
Limit simulator scratch regressions in the core benchmark suite to PhysX. OVPhysX-specific coverage belongs to its backend package and requires the optional runtime wheel.
3 participants
Description
This aggregate PR preserves the ancestry of the articulation reordering and selector stack so it can be reviewed and merged as one coherent change.
Merge order preserved in this branch:
int32/int64selector support inside Warp work kernelsProxyArrayfinder selectors and explicit legacy return modesThe P10 commits remain ancestors of P11. The source PRs should stay open and unchanged until this aggregate PR merges; they can then be closed as superseded.
Related overlap: #6674. This branch includes the OVPhysX Jacobian, mass-matrix, and gravity-compensation side folded into P9.
Key outcomes:
ProxyArrayselectors with zero-copy Torch and Warp views and per-asset LRU ownership.as_proxy=Falseretains them explicitly during migration.Asset performance before and after this PR
The table compares the PR merge-base (
a050ceade7) with the pushed PR head (00c8c25cb1) on an RTX 5090 using CUDA, 4,096 instances, 100 warmup calls, and 1,000 measured calls per operation. It contains only timed asset method/property calls; application and benchmark startup are excluded.Values are the equal-weight mean of the per-operation mean latencies for operations present on both revisions. The pre-PR
torch_tensormode and post-PRtorch_tensor_int32mode both generate CUDAtorch.int32selectors. Lower is better.API-method latency improved in 8 of 9 backend/asset pairs. Data-property latency improved in every comparable pair.
¹ The merge-base OVPhysX articulation data benchmark fails during
body_com_jacobian_wpreflight, while the post-PR suite passes. That data row is excluded instead of mixing unmatched measurements. The newint64and Warp selector modes are also excluded from this before/after table because they have no pre-PR equivalent; the post-PR all-mode matrix passes across all nine backend/asset pairs.Type of change
Verification
./isaaclab.sh -f: all hooks passed at the pushed HEAD.Full Kit-backed PhysX/Newton live asset suites were not completed locally because the available Kit environment timed out during startup; CI should provide that coverage.
Screenshots
Not applicable.
Checklist
./isaaclab.sh -fsource/<pkg>/changelog.d/for every touched packageCONTRIBUTORS.md