[DRAFT][Newton][Factory] Prereqs for NutThread support

[mzamoramora-nvidia]

Prereqs for Factory NutThread support on the Newton backend

Description

Aggregates the four IsaacLab changes required to support the Factory NutThread task on the Newton backend onto a single base branch. Each item is independently reviewable and would normally live in its own upstream PR; they are bundled here so the follow-up Factory branch (#5836 — Factory NutThread support on the Newton backend) can rebase against one base instead of juggling four moving targets while the upstream PRs land.

This branch is intentionally draft — once the constituent upstream PRs are merged, this branch will be dropped and the Factory PR will rebase onto develop directly. The bundle exists only to unblock end-to-end Factory training while review proceeds in parallel.

Relationship to the Factory NutThread PR

This PR is the base for #5836 — [Factory] NutThread support on the Newton backend. The commits here are also temporarily carried on the Factory branch so GitHub can render its diff; they belong to this PR's review surface.

Compare view (this PR's diff vs develop):
mzamoramora-nvidia/IsaacLab@develop...nut-thread-newton-prereqs

What this branch contains

Piece	Upstream PR	Author	Status
SDF mesh collision + hydroelastic shape configuration	#5228	@AntoineRichard	open
NewtonCfg.collision_decimation (mid-tick re-collide)	#5729	@mzamoramora-nvidia	open
RigidObjectData shape fix for kinematic single-body assets	n/a (new)	@mzamoramora-nvidia	this PR
NewtonManager.sanitize_world_state hook + deferred NaN-reset queue	n/a (new)	@mzamoramora-nvidia	this PR

Topology:

develop
  ├── merge PR #5228 (SDF / hydroelastic cfg)
  ├── merge PR #5729 (collision_decimation)
  ├── Newton: fix RigidObjectData shape for kinematic single-body assets
  └── newton_manager: add sanitize_world_state hook for NaN recovery

The two new commits sit on top of the vendored PRs because both depend on NewtonCollisionPipelineCfg being available in NewtonCfg.

Detail per piece

PR #5228 — SDF / hydroelastic configuration (vendored as-is)

Adds SDFCfg and surfaces the HydroelasticSDFCfg fields needed to drive Newton's pressure-based contact pipeline (moment_matching, buffer_mult_*, grid_size). No local modifications; merged verbatim so the Factory branch can configure sdf_hydroelastic_config=HydroelasticSDFCfg(...) against the same struct shape that will land upstream.

PR #5729 — NewtonCfg.collision_decimation (vendored as-is)

Adds the opt-in collision_decimation knob that re-invokes the Newton collision pipeline every N solver substeps inside a single physics tick. Default 0 preserves legacy behaviour. The Factory NutThread newton_sdf preset uses collision_decimation=1 (re-collide every substep); the newton (hydroelastic) preset uses collision_decimation=2. These values were chosen empirically — tighter decimation kept the SDF preset stable on this task; hydroelastic tolerated a looser setting.

RigidObjectData shape fix for kinematic single-body assets (new)

Fixes :class:~isaaclab_newton.assets.RigidObjectData crashing with IndexError: tuple index out of range when a fixed-base rigid object holds a single body. The is_fixed_base branch in _create_simulation_bindings indexed [:, 0, 0] assuming a 3D (count, links, 1) layout, but Newton returns a 2D (count, links) array when the view contains a single body. Dispatch on actual ndim so both multi-link fixed-base articulations and single-body kinematic rigid objects work. Also widens _sim_bind_body_com_vel_w to (num_instances, 1) so the no-velocity fallback matches the derive_body_acceleration_from_body_com_velocities kernel's 2D signature.

This is hit immediately when spawning the Factory bolt as a kinematic single-body RigidObjectCfg (the natural Newton-side representation since the bolt is fixed in the world); other tasks that use the same kinematic single-body pattern would hit it too.

NewtonManager.sanitize_world_state hook + deferred NaN-reset queue (new)

Adds two classmethods to NewtonManager:

• sanitize_world_state(env_ids) — zeroes per-world MJWarp solver internals (qacc_warmstart, qfrc_applied, xfrc_applied, qacc, cacc, cfrc_int, cfrc_ext) and Newton State velocity buffers for the given worlds. Discussed in newton-physics/newton#1266. The current list is intentionally defensive — we clear every MJWarp-internal buffer that could plausibly carry NaN state forward, and stopped tightening once recovery was reliable. Not all of these probably need to be reset; a follow-up should bisect which subset is actually load-bearing.
• flag_nan_envs(mask) / sanitize_pending_nan_envs() — a tiny classvar-backed pending-reset queue. Tasks call flag_nan_envs when they detect NaN (e.g. in _get_observations), and the env's reset path drains the queue via sanitize_pending_nan_envs once per reset. This defers the scrub to the natural episode boundary instead of clobbering live warm-starts mid-rollout — earlier eager-sanitize attempts regressed hydroelastic training, and deferring the scrub fixed the regression.

The buffers scrubbed are MJWarp-specific; the hook has only been exercised against SolverMuJoCo so far. Other Newton solvers may need a different set of buffers, or none at all — that's an open question for when a non-MJWarp task first needs NaN recovery.

Why bundle this with the prereqs rather than upstream-first: the API surface (two classmethods + one classvar) is small but the placement (which buffers to scrub, where to call the hook) is task-driven; the Factory NutThread env is the first task that actually exercises it. Once the contract is validated by landing the Factory PR, splitting this out into a standalone upstream PR is straightforward — it doesn't depend on anything else here.

Type of change

• New feature (non-breaking change which adds functionality) — net-new hooks and configclasses, no removal or signature change to existing public API.

Test plan

Check	Result
./isaaclab.sh -p -m pytest source/isaaclab_newton/test/physics/test_newton_manager_abstraction.py	39/39 pass (PR #5729 tests, unchanged)
./isaaclab.sh -p -m pytest source/isaaclab_newton/test/physics/test_sdf_config.py	14/14 pass (PR #5228 tests, unchanged)
./isaaclab.sh -f	clean (ruff, ruff-format, codespell, license, rst lint)
Factory smoke (drives the new code): Isaac-Factory-NutThread-Direct-v0 presets=newton	reaches 100% success at epoch 71 in a 200-epoch run on top of the dependent Factory PR

Standalone regression tests for the two new commits — single-body kinematic RigidObjectCfg round-trip and sanitize_world_state NaN-recovery — are open follow-ups to keep this PR scoped to vendoring + the minimal API additions the Factory branch depends on.

Checklist

• 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 (Sphinx docstrings on the new hook + classvar)
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works — vendored PRs ship with theirs; standalone tests for the two new commits are open follow-ups (see Test plan)
• I have added a changelog fragment under source/<pkg>/changelog.d/ for every touched package
• I have added my name to the CONTRIBUTORS.md or my name already exists there

[isaaclab-review-bot]

Code Review Summary

This PR bundles four prerequisite changes for Factory NutThread support on the Newton backend. The implementation is well-structured overall with proper configuration separation and comprehensive test coverage for the SDF configuration logic.

✅ Strengths

• Clean separation between configuration classes and implementation
• Proper handling of mid-tick re-collide skip on last substep
• Comprehensive unit tests for SDF configuration (14 tests)
• Good docstrings explaining the rationale for defensive buffer zeroing in sanitize_world_state

---

🔍 Findings

1. [Medium] Potential device mismatch in flag_nan_envs (newton_manager.py, ~line 1575)

When _nan_env_mask_pending_reset is first created via torch.zeros_like(mask), it inherits the device from the input mask. If a subsequent call passes a mask on a different device, the |= operation will fail.

Suggested fix:

if cls._nan_env_mask_pending_reset is None:
    cls._nan_env_mask_pending_reset = torch.zeros_like(mask)
elif cls._nan_env_mask_pending_reset.device != mask.device:
    mask = mask.to(cls._nan_env_mask_pending_reset.device)
cls._nan_env_mask_pending_reset |= mask

Alternatively, document that all masks must be on the same device as the first call.

---

2. [Medium] Defensive assertion for unexpected ndim (rigid_object_data.py, ~line 815)

The ndim >= 3 check correctly handles the 2D vs 3D dispatch, but an ndim of 1 or 4+ would silently take the "else" branch. Making the assumption explicit would fail fast on unexpected Newton API changes:

root_xforms = self._root_view.get_root_transforms(SimulationManager.get_state_0())
if root_xforms.ndim not in (2, 3):
    raise ValueError(f"Unexpected root_xforms ndim={root_xforms.ndim}, expected 2 or 3")
self._sim_bind_root_link_pose_w = root_xforms[:, 0, 0] if root_xforms.ndim == 3 else root_xforms[:, 0]

---

3. [Low/Suggestion] Add trace logging when sanitize_world_state is no-op (newton_manager.py, ~line 1511)

For debugging purposes, it can be helpful to know when sanitize_world_state is a no-op due to a non-MJWarp backend:

if mjw_data is None:
    logger.debug("sanitize_world_state: skipped (no mjw_data, likely non-MJWarp solver)")
    return

---

4. [Low/Suggestion] Duplicate regex compilation helper (newton_replicate.py, line 18)

The _compile_sdf_patterns function duplicates the same logic in newton_manager.py:_apply_sdf_config. Consider extracting to a shared utility in newton_collision_cfg.py to avoid drift between the two implementations.

---

📋 Test Coverage Notes

The PR acknowledges that standalone regression tests for the kinematic single-body fix and sanitize_world_state NaN recovery are follow-up items. Consider adding basic tests for the deferred NaN queue (flag_nan_envs/sanitize_pending_nan_envs) to prevent regressions in that coordination logic.