Fixes NaN after env reset by syncing body_q before collision detection (#5176)

# Description

This is an important fix for locomotion rough terrain in newton to work

- Fixed NaN explosions during Newton training caused by stale body_q
after env resets. Articulation write methods (write_root_pose_to_sim_*,
write_joint_position_to_sim_*) updated joint_q but never recomputed
body_q via forward kinematics. The collision pipeline read the old
body_q, generating contacts for the wrong position. The MuJoCo Warp
solver then computed constraint Jacobians from new qpos but old contacts
— a geometry mismatch that overflowed the solver gradient, producing
NaN.
- Added NewtonManager.invalidate_fk() flag set by articulation write
methods, checked in step() before collision detection to conditionally
run eval_fk.
- Added regression test that patches _simulate_physics_only to capture
body_q at collide() time and asserts consistency with joint_q. Fails
with 10m mismatch without the fix.

Fixes # (issue)

<!-- As a practice, it is recommended to open an issue to have
discussions on the proposed pull request.
This makes it easier for the community to keep track of what is being
developed or added, and if a given feature
is demanded by more than one party. -->

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- Bug fix (non-breaking change which fixes an issue)
- New feature (non-breaking change which adds functionality)
- Breaking change (existing functionality will not work without user
modification)
- Documentation update

## Screenshots

Please attach before and after screenshots of the change if applicable.

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

## Checklist

- [x] I have read and understood the [contribution
guidelines](https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html)
- [x] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [x] My changes generate no new warnings
- [x] I have added tests that prove my fix is effective or that my
feature works
- [x] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [x] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------

Co-authored-by: isaaclab-review-bot[bot] <isaaclab-review-bot[bot]@users.noreply.github.com>
Co-authored-by: Antoine RICHARD <antoiner@nvidia.com>

Author: 3 people

source/isaaclab_newton/config/extension.toml

@@ -1,7 +1,7 @@
 [package]
 
 # Note: Semantic Versioning is used: https://semver.org/
-version = "0.5.9"
+version = "0.5.10"
 
 # Description
 title = "Newton simulation interfaces for IsaacLab core package"

source/isaaclab_newton/docs/CHANGELOG.rst

@@ -1,6 +1,18 @@
 Changelog
 ---------
 
+0.5.10 (2026-04-05)
+~~~~~~~~~~~~~~~~~~~
+
+Fixed
+^^^^^
+
+* Fixed NaN after env reset caused by stale ``body_q`` in the collision
+  pipeline. Added :meth:`~isaaclab_newton.physics.NewtonManager.invalidate_fk`
+  so articulation write methods trigger ``eval_fk`` before the next
+  ``collide()``.
+
+
 0.5.9 (2026-03-16)
 ~~~~~~~~~~~~~~~~~~
 

source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py

@@ -495,6 +495,7 @@ def write_root_link_pose_to_sim_index(
         if self.data._root_state_w is not None:
             self.data._root_state_w.timestamp = -1.0
         self.data._fk_timestamp = -1.0  # Forces a kinematic update to get the latest body link poses.
+        SimulationManager.invalidate_fk()
         if self.data._body_com_pose_w is not None:
             self.data._body_com_pose_w.timestamp = -1.0
         if self.data._body_state_w is not None:
@@ -550,6 +551,7 @@ def write_root_link_pose_to_sim_mask(
         if self.data._root_state_w is not None:
             self.data._root_state_w.timestamp = -1.0
         self.data._fk_timestamp = -1.0  # Forces a kinematic update to get the latest body link poses.
+        SimulationManager.invalidate_fk()
         if self.data._body_com_pose_w is not None:
             self.data._body_com_pose_w.timestamp = -1.0
         if self.data._body_state_w is not None:
@@ -614,6 +616,7 @@ def write_root_com_pose_to_sim_index(
         if self.data._root_state_w is not None:
             self.data._root_state_w.timestamp = -1.0
         self.data._fk_timestamp = -1.0  # Forces a kinematic update to get the latest body link poses.
+        SimulationManager.invalidate_fk()
         if self.data._body_com_pose_w is not None:
             self.data._body_com_pose_w.timestamp = -1.0
         if self.data._body_state_w is not None:
@@ -674,6 +677,7 @@ def write_root_com_pose_to_sim_mask(
         if self.data._root_state_w is not None:
             self.data._root_state_w.timestamp = -1.0
         self.data._fk_timestamp = -1.0  # Forces a kinematic update to get the latest body link poses.
+        SimulationManager.invalidate_fk()
         if self.data._body_com_pose_w is not None:
             self.data._body_com_pose_w.timestamp = -1.0
         if self.data._body_state_w is not None:
@@ -1012,6 +1016,7 @@ def write_joint_position_to_sim_index(
         )
         # Invalidate FK timestamp so body poses are recomputed on next access.
         self.data._fk_timestamp = -1.0
+        SimulationManager.invalidate_fk()
         # Need to invalidate the buffer to trigger the update with the new root pose.
         # Only invalidate if the buffer has been accessed (not None).
         if self.data._body_link_vel_w is not None:
@@ -1066,6 +1071,7 @@ def write_joint_position_to_sim_mask(
         )
         # Invalidate FK timestamp so body poses are recomputed on next access.
         self.data._fk_timestamp = -1.0
+        SimulationManager.invalidate_fk()
         # Need to invalidate the buffer to trigger the update with the new root pose.
         # Only invalidate if the buffer has been accessed (not None).
         if self.data._body_link_vel_w is not None:

source/isaaclab_newton/isaaclab_newton/physics/newton_manager.py

@@ -93,6 +93,7 @@ class NewtonManager(PhysicsManager):
     _collision_pipeline = None
     _newton_contact_sensors: dict = {}  # Maps sensor_key to NewtonContactSensor
     _report_contacts: bool = False
+    _fk_dirty: bool = False
 
     # CUDA graphing
     _graph = None
@@ -158,6 +159,7 @@ def reset(cls, soft: bool = False) -> None:
     def forward(cls) -> None:
         """Update articulation kinematics without stepping physics."""
         eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None)
+        cls._fk_dirty = False
 
     @classmethod
     def pre_render(cls) -> None:
@@ -298,6 +300,13 @@ def step(cls) -> None:
             else:
                 logger.warning("Newton deferred CUDA graph capture failed; using eager execution")
 
+        # Ensure body_q is up-to-date before collision detection.
+        # After env resets, joint_q is written but body_q (used by
+        # broadphase/narrowphase) is stale until FK runs.
+        if cls._fk_dirty and cls._needs_collision_pipeline:
+            eval_fk(cls._model, cls._state_0.joint_q, cls._state_0.joint_qd, cls._state_0, None)
+        cls._fk_dirty = False
+
         # Step simulation (graphed or not; _graph is None when capture is disabled or failed)
         if cfg is not None and cfg.use_cuda_graph and cls._graph is not None and "cuda" in device:  # type: ignore[union-attr]
             wp.capture_launch(cls._graph)
@@ -358,6 +367,7 @@ def clear(cls):
         cls._collision_pipeline = None
         cls._newton_contact_sensors = {}
         cls._report_contacts = False
+        cls._fk_dirty = False
         cls._graph = None
         cls._graph_capture_pending = False
         cls._newton_stage_path = None
@@ -377,6 +387,16 @@ def add_model_change(cls, change: SolverNotifyFlags) -> None:
         """Register a model change to notify the solver."""
         cls._model_changes.add(change)
 
+    @classmethod
+    def invalidate_fk(cls) -> None:
+        """Mark forward kinematics as needing recomputation.
+
+        Called by articulation write methods that modify ``joint_q`` or root
+        transforms.  The flag is checked in :meth:`step` before collision
+        detection to ensure ``body_q`` is up-to-date.
+        """
+        cls._fk_dirty = True
+
     @classmethod
     def start_simulation(cls) -> None:
         """Start simulation by finalizing model and initializing state.

source/isaaclab_newton/test/assets/test_articulation.py

@@ -2406,5 +2406,86 @@ def test_write_joint_frictions_to_sim(sim, num_articulations, device, add_ground
     assert torch.allclose(joint_friction_coeff_sim_2, friction_2)
 
 
+@pytest.mark.parametrize("num_articulations", [2])
+@pytest.mark.parametrize("device", ["cuda:0"])
+@pytest.mark.parametrize("articulation_type", ["anymal"])
+def test_body_q_consistent_after_root_write(num_articulations, device, articulation_type):
+    """Test that body_q is fresh when collide() runs after a root pose write.
+
+    Regression test for a NaN bug where collide() used stale body_q after env
+    reset because eval_fk was not called between write_root_pose and collide.
+
+    Uses ``use_mujoco_contacts=False`` so the Newton collision pipeline is
+    active, then patches ``_simulate_physics_only`` to capture body_q at
+    the moment collide() is called and asserts it matches joint_q.
+    """
+    from unittest.mock import patch
+
+    sim_cfg = SimulationCfg(
+        dt=1 / 200,
+        physics=NewtonCfg(
+            solver_cfg=MJWarpSolverCfg(
+                njmax=70,
+                nconmax=70,
+                integrator="implicitfast",
+                use_mujoco_contacts=False,
+            ),
+            num_substeps=1,
+            use_cuda_graph=False,
+        ),
+    )
+    with build_simulation_context(sim_cfg=sim_cfg, device=device) as sim:
+        sim._app_control_on_stop_handle = None
+        articulation_cfg = generate_articulation_cfg(articulation_type=articulation_type)
+        articulation, env_pos = generate_articulation(articulation_cfg, num_articulations, device)
+
+        sim.reset()
+
+        model = SimulationManager.get_model()
+        jc_starts = model.joint_coord_world_start.numpy()
+        body_starts = model.body_world_start.numpy()
+
+        for _ in range(5):
+            sim.step()
+            articulation.update(sim.cfg.dt)
+
+        # Teleport env 0 by 10m (simulating a reset)
+        new_pose = wp.to_torch(articulation.data.default_root_pose).clone()
+        new_pose[0, 0] += 10.0
+        new_pose[0, 1] += 5.0
+        articulation.write_root_pose_to_sim_index(
+            root_pose=new_pose[0:1],
+            env_ids=torch.tensor([0], device=device, dtype=torch.int32),
+        )
+
+        # Patch _simulate_physics_only to capture body_q before collide runs
+        captured = {}
+        original_simulate = SimulationManager._simulate_physics_only.__func__
+
+        @classmethod  # type: ignore[misc]
+        def _patched_simulate(cls):
+            if cls._needs_collision_pipeline:
+                bq = wp.to_torch(cls._state_0.body_q)
+                jq = wp.to_torch(cls._state_0.joint_q)
+                b0 = int(body_starts[0])
+                jc0 = int(jc_starts[0])
+                captured["bq_root"] = bq[b0, :3].clone()
+                captured["jq_root"] = jq[jc0 : jc0 + 3].clone()
+            original_simulate(cls)
+
+        with patch.object(SimulationManager, "_simulate_physics_only", _patched_simulate):
+            sim.step()
+        articulation.update(sim.cfg.dt)
+
+        assert captured, "collision pipeline did not run — _needs_collision_pipeline is False"
+
+        bq_root = captured["bq_root"]
+        jq_root = captured["jq_root"]
+        diff = (jq_root - bq_root).abs().max().item()
+        assert diff < 0.01, (
+            f"body_q was stale when collide() ran: diff={diff:.4f}m, jq={jq_root.tolist()}, bq={bq_root.tolist()}"
+        )
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "--maxfail=1"])