Articulation Reordering Part 11: Cache proxy selectors - #6740
Articulation Reordering Part 11: Cache proxy selectors#6740AntoineRichard wants to merge 55 commits into
Conversation
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.
Keep finder proxies in Warp-backed form through dispatchers, rigid writers, and repeated joint actions. Preserve legacy subset finder indices while caching asset-global writer selectors.
Surface immediate failures for registered selector modes instead of reporting skips. Verify successful writer summaries retain all six registered modes in both backends.
Record explicit proxy and legacy finder modes for each backend. Enforce lifetime, view-access, subset-index, and changelog migration guidance with documentation contract tests.
Validate only the selector modes requested by CLI orchestration, preserving intentional partial summaries while making all-mode gaps fatal. Surface missing registered writer methods in both backends.
Advertise ProxyArray on deprecated collection object selectors without changing write_data_to_sim. Tailor articulation joint and tendon examples to their matching writers and selector names.
|
Too many files changed for review. ( Bypass the limit by tagging |
github-actions
Bot
added
documentation
![isaaclab-review-bot[bot]](https://avatars.githubusercontent.com/u/157846462?s=60&v=4){kind=small}
There was a problem hiding this comment.
Isaac Lab Review Bot
Summary
This change adds an opt-in cached selector path (as_proxy=True) to asset finder APIs across the core, PhysX, Newton, and OVPhysX packages, backed by a per-asset 128-entry LRU of device-local wp.int32 ProxyArray selectors, and threads those selectors through actuators, action terms, and backend index-kernel dispatchers. The deprecation shape is well handled: omitted as_proxy keeps the legacy return with a DeprecationWarning, as_proxy=False keeps it silently, and call sites throughout the repo were updated explicitly rather than left to warn. Changelog fragments are present for each touched package and the kernel/selector docstrings follow the Google-style and SI-unit conventions. The main remaining concerns are the mutability contract on cached selectors (documented but not enforced), a small regression in actuator __str__ for partial slices, one documentation command that contradicts the repo's tooling guidance, and several new test files that assert on source text or on transient changelog fragments rather than on runtime behavior.
Architecture impact
The cached-selector layer is introduced at the AssetBase boundary (_AssetSelectorCache plus _resolve_finder_indices), which keeps the deprecation/return-mode policy in one place and lets all three backends and the mocks share it — a good factoring. IndexKernelDispatcher similarly centralizes int32/int64 kernel specialization instead of duplicating overloads per backend. The cost is a new invariant that spreads across many call sites: finder results may now be shared, aliased objects, so every downstream consumer must branch on slice vs ProxyArray (visible in the repeated joint_ids = ... if isinstance(..., slice) else ....torch pattern in the action terms). That pattern is acceptable during the transition but is worth consolidating into a helper before it multiplies further.
Test coverage
Coverage of the new surface is broad: cache identity and LRU eviction, per-backend finder return modes for articulations/rigid objects/collections, resolver unwrapping without materializing Torch, empty-array .torch views, and index-kernel dtype dispatch. Weak spots: several new suites (test_articulation_subset_proxy_finders.py, test_joint_action_proxy_selectors.py, test_finder_proxy_documentation.py) validate parsed source text or changelog fragments rather than the imported objects, so they cannot fail if an action term or backend finder stops requesting proxies; the empty-array DLPack branch and the device-local allocation guarantee are exercised only on CPU; and one negative dtype test matches a bare RuntimeError, which many unrelated Warp failures also raise.
Implementation verdict
Minor fixes needed. Posted 7 actionable findings inline.
The PR exceeded the automated context budget, so some file content was truncated.
Automated comment-only review; human maintainers own approval decisions.
| self._capacity = capacity | ||
| self._entries: OrderedDict[tuple[str, tuple[int, ...]], ProxyArray] = OrderedDict() | ||
|
|
||
| def get(self, domain: str, indices: Sequence[int], device: str) -> ProxyArray: |
There was a problem hiding this comment.
🟡 Warning — Cached selectors hand out shared, writable storage
_AssetSelectorCache.get returns the same ProxyArray instance (and the same underlying wp.int32 allocation) for every equivalent index sequence, while both .warp and .torch stay writable. Legacy finder results were independently-owned values, so existing code that mutates a finder result in place — e.g. an in-place Torch op on joint_ids — would now silently corrupt the cache and make unrelated later callers target the wrong joints or bodies, with no error at the point of misuse. The docstrings state callers must treat the proxy as immutable, but nothing enforces it. Consider marking the cached Warp array (and the derived Torch view) read-only at the boundary, or returning an independent copy whenever writable access is exposed.
| "source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object/rigid_object.py", | ||
| "source/isaaclab_ovphysx/isaaclab_ovphysx/assets/rigid_object_collection/rigid_object_collection.py", | ||
| ) | ||
| _BACKEND_FRAGMENTS = ( |
There was a problem hiding this comment.
🟡 Warning — Test pins the existence of transient changelog.d fragments
_BACKEND_FRAGMENTS asserts on source/<pkg>/changelog.d/articulation-reordering-p11.rst. Those fragments are consumed and removed by the nightly changelog compile, so once that runs this test fails permanently on develop even though the finder behavior is unchanged. It also duplicates validation the changelog tooling already performs and proves nothing about the API this file is otherwise checking. Drop the fragment assertions and keep only the docstring/signature checks over _FINDER_FILES.
|
|
||
| ./isaaclab.sh train --rl_library rsl_rl \ | ||
| --task Isaac-Velocity-Flat-AnymalD \ | ||
| --num_envs 4096 \ |
There was a problem hiding this comment.
🔵 Suggestion — New guide invokes ./isaaclab.sh for training and playback
The quick-transfer examples use ./isaaclab.sh train ... and ./isaaclab.sh play .... Per the repository tooling guidance, ./isaaclab.sh is reserved for advanced installation workflows (-i <args>), and routine CLI workflows should go through uv run isaaclab. Please update the training and playback commands in this new page so readers land on the supported invocation path.
| --num_envs 4096 \ | |
| uv run isaaclab train --rl_library rsl_rl \ |
| # resolve joint indices for printing | ||
| joint_indices = self.joint_indices | ||
| if joint_indices == slice(None): | ||
| if isinstance(joint_indices, slice): |
There was a problem hiding this comment.
🔵 Suggestion — Partial joint slices now print as the full local joint range
The guard changed from if joint_indices == slice(None): to if isinstance(joint_indices, slice):, so any slice — including a partial one such as slice(2, 5) — is replaced by list(range(self.num_joints)). For an actuator built with a partial slice, __str__ then reports joint IDs [0, 1, ...] instead of the actual articulation joints, producing misleading resolution diagnostics. Keep the type guard (needed to avoid ambiguous tensor comparison) but restrict the full-range expansion to the slice(None) sentinel.
| if isinstance(joint_indices, slice): | |
| if isinstance(joint_indices, slice) and joint_indices == slice(None): |
| "joint": _REPO_ROOT / "source/isaaclab/isaaclab/envs/mdp/actions/joint_actions.py", | ||
| "limits": _REPO_ROOT / "source/isaaclab/isaaclab/envs/mdp/actions/joint_actions_to_limits.py", | ||
| } | ||
| _WRITER_PATHS = ( |
There was a problem hiding this comment.
🔵 Suggestion — AST-exec test cannot detect action terms dropping as_proxy=True
This suite parses the action modules, extracts writer methods with ast, and execs them against SimpleNamespace stubs. Because the bodies run detached from their class, the test never observes what __init__ actually stored: if JointAction.__init__ or JointPositionToLimitsAction.__init__ reverted to a list/tensor selector, every assertion here would still pass — which is exactly the regression this PR is guarding. Add at least one behavioral test that constructs the action term against MockArticulation and asserts isinstance(term._joint_ids, ProxyArray) and that the writer receives that same object.
| """The public Wrench raw symbol accepts direct launches only with int32 selectors.""" | ||
| env_ids = wp.array([0], dtype=index_dtype, device="cpu") | ||
| buffers = [wp.zeros((1, 1), dtype=wp.vec3f, device="cpu") for _ in range(7)] | ||
| with pytest.raises(RuntimeError): |
There was a problem hiding this comment.
🔵 Suggestion — Negative dtype test accepts any RuntimeError
test_public_wrench_reset_kernel_rejects_non_int32_direct_launch only asserts that wp.launch raises a bare RuntimeError. Warp raises RuntimeError for argument-count, shape, and device mismatches too, so this would keep passing if the kernel signature changed arity or the seven vec3f buffers stopped matching — it does not actually pin the int32-only selector contract. Match on the error text so the assertion is sensitive to the dtype specifically.
| with pytest.raises(RuntimeError): | |
| with pytest.raises(RuntimeError, match="int32"): |
| t2 = ta.torch | ||
| assert t1 is t2 | ||
|
|
||
| def test_empty_cpu_torch_returns_cached_tensor(self): |
There was a problem hiding this comment.
🔵 Suggestion — New empty-array DLPack branch is only covered on CPU
ProxyArray.torch now takes a separate torch.from_dlpack path when self._warp.size == 0, including manual requires_grad/grad propagation. Both new tests hardcode device="cpu" even though this test class already has a device fixture parametrized over CPU and CUDA. Zero-size DLPack export and grad-tensor construction are exactly the places where the two backends can diverge, and cached finders can legitimately return empty selectors on CUDA assets. Use the existing fixture and build the warp arrays on device.
| def test_empty_cpu_torch_returns_cached_tensor(self): | |
| def test_empty_cpu_torch_returns_cached_tensor(self, device): |
1 participant
Description
Adds an opt-in cached selector path to asset
find_*APIs. Withas_proxy=True, articulation, rigid-object, and rigid-object-collectionfinders return a per-asset cached
ProxyArraybacked by device-localwp.int32storage. The same allocation can be consumed directly by Warpwriters or accessed as a zero-copy Torch view.
The existing finder result types remain available during the deprecation
period: omitted
as_proxyretains the legacy result and warns, whileas_proxy=Falseexplicitly retains it without warning. The cache is aper-asset 128-entry LRU and is cleared when the asset is invalidated.
PhysX, Newton, and OVPhysX selector boundaries consume the cached Warp
allocation without materializing Torch. Repeated actuator and action-term
writer paths now retain proxy selectors, and the articulation benchmark
grid covers list, Torch int32/int64, Warp int32/int64, and cached proxy
representations.
This is stacked after #6652.
Note
The articulation reordering parts target
develop, so until the precedingparts merge this PR's diff includes them. The true P11 delta is
this compare.
Type of change
Validation
./isaaclab.sh -fdocumentation suites: 186 passed.
findings.
GPU-backed method-grid timings were not rerun for P11 because the shared GPU
was occupied; the benchmark registration, synchronization, completeness, and
reporting contracts are covered by the focused tests.
Checklist
./isaaclab.sh -fCONTRIBUTORS.mdDocumentation screenshots are not applicable.