Skip to content

Unify Warp math, memory operations, and sensor capabilities - #8820

Draft
ooctipus wants to merge 4 commits into
isaac-sim:mainfrom
ooctipus:refactor/shared-math-kernels
Draft

ooctipus wants to merge 4 commits into
isaac-sim:mainfrom
ooctipus:refactor/shared-math-kernels

Conversation

@ooctipus

@ooctipus ooctipus commented Oct 26, 2026 •

Copy link
Copy Markdown
Collaborator

Description

Frame math, packing, selected copies, and PVA storage are repeated across the OV, PhysX, Newton, and experimental packages. This change gives each shared concept one owner and generates concrete Warp kernels before launch. Callers resolve attribute frequencies into arrays and selectors; backends bind native inputs and declare the outputs they support.

Ownership and composition

  • isaaclab.utils.warp.math owns quaternion/frame operations, rigid-body transport, interval scaling, and COM-to-origin Jacobian shifts. memory owns value-preserving concatenation, splitting, views, clearing, selected copy/fill, and selection masks. It contains no body/joint/velocity policy, frequency enum, or fixed 13-component layout. The existing rigid-object data contract owns the public state type.
  • Callers supply element types, rank, strides and index/mask selections. Cached memory factories use isolated Warp modules, so importing another layout does not invalidate compiled variants. Newton retains its existing AttributeFrequency/view layouts; unequal joint-coordinate and joint-DoF layouts still require explicit conversion. Mathematical binary variants remain in their owner's module. Wrench and articulation-ordering specializations are declared before launch.
  • 86 backend copy/fill launch sites now use one generated writer, preserving Cartesian selection, mixed int32/int64 selectors, full/compact sources, partial writes, and optional simulator row-ID outputs. Another root copy uses its rank-1 variant and retains cached launching. 28 OV scalar property paths no longer allocate broadcast Torch tensors. Native solver updates, inverse mass/inertia calculations, COM conversion, fused friction updates, and staging remain backend-owned.
  • Shared IMU, PVA, and frame-transformer equations stay inside fused sensor kernels. Backend packages retain native gathering, input layouts, and solver behavior.
  • Core PvaData owns standard outputs. Pose/position/quaternion share one allocation; each linear/angular pair shares a spatial vector. Getters return cached Warp/Torch views and launch no packing kernels. The abstract data shell, backend data factory/classes, and forwarding exports are removed.
  • PvaCfg.data_types requests outputs; each sensor backend declares supported_data_types. Pose is required; velocity, acceleration, and projected gravity are independently optional. Eight import-time PVA variants reuse one equation writer. Unsupported requests fail before callbacks/resources are created. Pose-only Newton does not request body_qdd; acceleration-only still gathers angular velocity for transport.
  • Additional backend measurements can use a local data subclass that passes its standard subset to PvaData and owns extra storage, updates, and resets. Core does not impose a universal feature set or extension registry.
  • Mask resolution has one core owner and preserves integer widths, strides, negative indices, and slice semantics. Preallocated same-device selectors are capturable. Eager temporary GPU index uploads synchronize their stream before releasing storage. Dead joint-action mask state is removed.

All in-tree callers migrate directly without compatibility aliases. Architecture gates reject retired owners/imports, duplicate definitions, reversed dependencies, and late kernel specialization. The Torch math/autograd API remains unchanged.

This also fixes the experimental reorientation task's xyzw quaternion indexing: identity-to-identity previously returned π and an X-axis half turn returned 0; the shared shortest-angle function returns 0 and π. Frame-transformer source outputs have one writer per environment. Wrench exclusive-cell accumulation avoids accidental Warp atomics.

API and lifecycle changes

Import PvaData from isaaclab.sensors; backend PvaData and BasePvaData imports are removed. Import retired backend kernel helpers from their core math, memory, sensor, or wrench owner; import resolve_1d_mask from core memory.

PVA component views now alias their grouped storage and may be noncontiguous. Reinitialization creates new views; retained old views retain old storage. Disabled outputs raise NotImplementedError, while enabled outputs return None before allocation. Defaults still enable all four standard groups. Debug visualization requires acceleration.

PVA initialization/reset/STOP now consistently produce identity pose, zero velocity/acceleration, and projected gravity along -Z until the next sample. This intentionally replaces Newton's zero-gravity initialization and zero-quaternion STOP state and clears stale OV/PhysX measurements on STOP. STOP also clears pending lazy sampling before native handles are released. Migration documentation and breaking-change markers cover these changes.

The cumulative change touches 145 files and removes 2,189 net lines. The frequency-independent memory revision removes 886 net production lines; added tests and documentation reduce its overall net reduction to 688 lines. There are no unmerged PR dependencies. The branch starts from main at a9fb50e5388724a680bbc5c1f836917271605a99; its working branch lives only in ooctipus/IsaacLab.

Type of change

  • Architectural refactor of shared Warp operations, generation, and sensor ownership
  • Bug fixes and documented API/lifecycle changes
  • Documentation and structural/numerical regression coverage

Validation

  • 552 tests passed across 13 isolated per-file runs for this revision: memory, architecture, integer specialization, shared sensors, articulation ordering, shared rigid-object/articulation contracts, and real Newton asset/tendon/PVA/frame-transformer/IMU/joint-wrench integration. These runs included 298 skips, 99 deselections and one expected failure. The contract tests use mocked native views; they are not full simulator validation.
  • Final memory/architecture checks passed after dtype validation and rank-1 generation cleanup. Coverage includes unfamiliar float64 component layouts, strided views, mixed-width Cartesian indices, duplicate IDs, full/compact sources, scalar fills, empty selections, untouched cells, and CUDA graph replay with changed data and selectors. The earlier 209-test validation remains evidence for unchanged math, lifecycle and experimental-task changes.
  • Centralized existence/ownership audit checked all 147 changed paths without rename folding. An independent AST audit checked argument order, destinations and dimensions at all 86 migrated launches. Deliberately adding a domain identifier to memory and corrupting selected writes failed their structural/numerical gates; both mutations were restored. Retired owners, aliases, reversed dependencies, duplicate state-type ownership and late specialization are forbidden.
  • All 19 live generated memory bindings, including their integer-width variants, and six remaining backend kernel modules compiled on CUDA.
  • Formatting, lint, spelling, changelog checks, and git diff --check passed. Clean Sphinx build passed with warnings treated as errors. Documentation imports real Warp type metadata instead of mocking it; specialization does not require launching GPU kernels.

Measured tradeoffs

On RTX 5090 / Warp 1.17.0, full PVA output storage falls from 116 to 88 bytes per environment (24.1%); pose-only output uses 28 bytes. This excludes backend input/staging storage.

Against the preceding PR commit, a warm PVA update plus one pose read is 22–29% faster at 4,096/65,536 environments because getters no longer repack pose. With no pose reads, the update kernel is 10–24% slower (about 0.3–1.1 µs). These are medians of 21 alternating trials, each replaying 200 steps in a CUDA graph, with numerical parity checked before timing. They exclude allocation, compilation, Python dispatch, and downstream consumers; no overall simulation/training speedup is claimed.

Against ff8e773bc9, the final memory revision measured packing within 1.2%, scalar fills 4–6% faster, and rank-3 trailing-component copies 52–54% faster. The general indexed copy is 0.09–0.13 µs slower (7–10%); the strided root copy is 0.18 µs slower on GPU, with cached eager cost changing from 2.87 to 2.90 µs. These ten cases use the same 21-trial/200-operation CUDA-graph method; eager root-copy timing uses seven trials of 2,000 calls. Architectural reuse is the primary benefit, not a universal speedup.

Remaining validation: OVPhysX is unavailable in this environment, and the PhysX recorded-launch suite could not bootstrap because optional Isaac Sim Kit is absent. Full OV/PhysX simulator suites, actual PhysX wrench delivery, full training, and distributed runs remain unrun. Focused runs emitted dependency warnings, so the warning-free checklist remains unchecked. The PR remains draft for review and broader simulator validation.

Checklist

  • I have read and understood the contribution guidelines.
  • I have run the pre-commit checks with uv run isaaclab --format on the changed files.
  • I have made corresponding changes to the documentation.
  • My changes generate no new warnings.
  • I have validated changed behavior with focused coverage and applied the test-audit criteria.
  • I have added changelog fragments for every touched package.
  • My name already exists in CONTRIBUTORS.md.

@github-actions github-actions Bot added documentation Improvements or additions to documentation isaac-lab Related to Isaac Lab team labels Oct 26, 2026
@ooctipus ooctipus changed the title Centralize shared Warp math and memory operations Unify Warp math, memory, and PVA capabilities Oct 26, 2026
@ooctipus ooctipus changed the title Unify Warp math, memory, and PVA capabilities Unify Warp math, memory operations, and sensor capabilities Oct 26, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation isaac-lab Related to Isaac Lab team

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant