fix(newton): resolve environment-vs-world index in the FK/reset-mask kernels

[pascal-roth]

Description

NewtonManager.invalidate_fk and invalidate_body_state (in isaaclab_newton) mark environments
as needing FK recomputation and solver reset by writing world_mask/fk_mask through four Warp
kernels (_or_reset_masks_from_mask, _scatter_reset_masks_from_ids,
_or_world_reset_mask_from_mask, _scatter_world_reset_mask_from_ids) that all conflated the
environment index with the world index.

That is correct for the replicated builder (world_count == num_envs, one world per environment)
but wrong for the flat (non-replicated) builder used when replicate_physics=False, where every
environment shares a single world 0 and the environment index is carried by the articulation
("arti") axis instead. Treating an environment index as a world index there wrote out of bounds
into world_mask (sized world_count == 1), producing an "illegal memory access" CUDA error on
every env.reset() call for any non-replicated scene with more than one environment.

The fix resolves the builder convention once at launch time (world_count == 1) and passes it
into each kernel so it indexes the correct axis.

This is the third bug in the same family (environment-index-vs-world-index confusion in Newton's
flat/non-replicated builder), alongside #7672 and #7673.

Reproduction / verification: reproduced against a custom multi-environment task with
replicate_physics=False: env.reset() crashed with a CUDA illegal-memory-access error before the
fix, confirmed synchronously with CUDA_LAUNCH_BLOCKING=1, which pinpointed the crash to
_scatter_reset_masks_from_ids. After the fix, the same task runs a full episode and records video
with no errors.

Fixes # (issue)

Type of change

• Bug fix (non-breaking change which fixes an issue)

Release backport

• Backport this pull request to the active release branch after it merges into develop

Screenshots

N/A — backend kernel fix, no visual output.

Checklist

Docker and GPU tests run on demand. Push the commits you want tested, then
comment run-ci on the pull request.

• I have read and understood the contribution guidelines
• I have run the pre-commit checks with ./isaaclab.sh --format
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• I have added a changelog fragment under source/<pkg>/changelog.d/ for every touched package (do not edit CHANGELOG.rst or bump extension.toml — CI handles that)
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[greptile-apps]

Confidence Score: 3/5

The PR is not safe to merge until flat body-state invalidation recognizes the actual mask layout and FK invalidation distinguishes flat builders from replicated single-environment scenes.

Findings

1.  Flat builder is never detected ▶
2.  Shape check misclassifies replicated scenes ▶

Summary

• Adds launch-time flat-builder selection to four reset-mask kernels.
• Maps flat-builder environment invalidations to shared world 0.
• Adds a changelog entry describing the CUDA out-of-bounds fix.
• The body-state path currently derives builder mode from the length of a mask containing an extra global slot, so that portion of the fix remains ineffective.
• FK invalidation also uses an ambiguous shape heuristic that breaks replicated single-environment scenes containing multiple articulations.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Environment invalidation] --> B{Builder layout}
  B -->|Replicated| C[world = env_id]
  B -->|Flat| D[world = 0]
  C --> E[Set world reset mask]
  D --> E
  A --> F[Set affected FK mask entries]

Loading

[greptile-apps]

Flat builder is never detected

In a flat multi-environment model, _world_reset_mask has world_count + 1 entries because it includes the reserved global slot. With world_count == 1, its length is therefore 2, so this condition never identifies the flat builder. invalidate_body_state consequently uses environment IDs as world indices: environment 1 dirties the reserved global slot, while larger IDs write out of bounds and can cause the same CUDA illegal-memory-access failure this change is intended to fix.

[greptile-apps]

Shape check misclassifies replicated scenes

A replicated scene with one environment and multiple articulations also has an articulation mapping shaped (1, >1), so this check classifies it as flat. The mask kernel then uses each articulation index to access env_mask, whose documented length is one. Every articulation after the first therefore reads out of bounds and can cause an illegal CUDA memory access in this valid replicated configuration.

[isaaclab-review-bot]

Isaac Lab Review Bot

The reset-mask kernels now map flat-builder environments to world 0, addressing the reported out-of-bounds world-mask writes. However, the sparse FK path marks every articulation in flat mode rather than only selected environments, and the shape-based flat-builder heuristic can misclassify a replicated single-environment view with multiple articulations and access env_mask out of bounds.

• Design and architecture: Passing the builder convention into the kernels is a focused design, but the convention is inferred using divergent shape heuristics. It should come from one authoritative topology condition that distinguishes a multi-environment flat builder from a replicated single-environment model.
• API: The public invalidate_fk and invalidate_body_state signatures and documented articulation_ids shape remain unchanged. Only private kernel signatures change, and the user-visible fix has an appropriate changelog fragment.
• Implementation: The world-mask indexing is corrected for multi-environment flat builders. Before merge, the sparse FK scatter must select only the articulation corresponding to each requested environment in flat mode, and flat-builder detection must avoid treating a replicated one-world view with multiple articulations as flat.

Minor fixes needed. Posted 2 actionable findings inline.

Automated review; human maintainers own approval decisions.

[isaaclab-review-bot]

🟡 Warning · Implementation — Flat scatter path marks FK for all environments

In flat mode world collapses to 0, yet the launch still spans articulation_ids.shape[1], which per this PR's own docstring is the environment axis, and fk_mask[articulation_ids[0, arti]] is set for every arti. A sparse invalidate_fk(env_ids=...) therefore marks FK for every environment, defeating the selectivity of the index path and risking FK recomputation over environments whose body state was written directly. In flat mode, index articulation_ids[0, env_id] with the launch dimensioned over env_ids only.

[isaaclab-review-bot]

🟡 Warning · Implementation — Flat-builder detection diverges from documented world_count

The kernel docstring defines the flag as world_count == 1, and states the replicated builder may have any arti count. This heuristic classifies any articulation_ids shaped (1, >1) as flat, so a replicated single-environment view with several articulations per world reads env_mask[arti] past its (num_envs,) = (1,) end — the same class of out-of-bounds access this PR fixes. Derive the flag from cls._model.world_count (already used in the fallback at line 1529), matching the stricter test used in invalidate_body_state.

[hujc7]

Agent review: Mapping environments onto their world is necessary, but one of the four kernels over-marks in the flat layout and the guard for one pair can never fire. Details inline; please add regression tests for sparse env_ids, mask-based selection, and >1 selected articulation per environment.

[hujc7]

[P0] Agent found bug: Always 0, so the invalidate_body_state half is dead. _world_reset_mask is allocated world_count + 1 long (:1685) — a single-world model has length 2, so shape[0] == 1 is unreachable; and _num_envs is None/1 on this path (measured), so the second conjunct fails independently. With _num_envs recorded, world_count == 1 and num_envs > 1 evaluates True on the same scene.

[hujc7]

[P0] Agent found bug: Measured on a flat Cartpole × 4 envs with this PR: invalidate_fk(env_ids=[1]) → fk_reset_mask=[1, 1, 1, 1] — every articulation flagged for one requested environment. The launch still spans (len(env_ids), articulation_ids.shape[1]), and in the flat layout arti is the environment axis, so each requested env marks the whole row. The mask path is correct (env_mask=[0,1,0,0] → [0, 1, 0, 0]). Setting the shared world's flag is right; the FK flags must select only the requested environments' articulations — either an explicit env→articulation mapping, or the narrower contract stated and enforced (one instance per env, in env order), in which case arti == env_id is the filter.

[hujc7]

[P1] Agent review: From reading: shape[0] == 1 and shape[1] > 1 cannot distinguish "N flat environments" (measured here: articulation_ids.shape=(1, 4)) from "one environment whose view selects M articulations" — the latter would read env_mask[arti] past a length-1 mask. Same root cause as above: the flat/replicated distinction needs to be recorded, not inferred from a shape.