Fix actuator gain snapshot corruption for multi-env floating-base articulations

source/isaaclab_newton/changelog.d/fix-actuator-gain-env-stride.rst

@@ -0,0 +1,12 @@
+Fixed
+^^^^^
+
+* Fixed the initial actuator gain snapshot used by
+  :func:`~isaaclab.envs.mdp.events.randomize_actuator_gains` corrupting (or
+  crashing) for multi-environment floating-base articulations with Newton
+  actuators. The per-environment stride of the actuator DOF indices was
+  decoded with the articulation-local joint count instead of the whole
+  model's per-environment DOF count, so on a floating base the free-root DOFs
+  shifted every environment past the first to the wrong (or out-of-bounds)
+  snapshot rows, corrupting the ``stiffness`` / ``damping`` randomization
+  baseline.

source/isaaclab_newton/isaaclab_newton/actuators/adapter.py

@@ -237,6 +237,7 @@ def build_newton_actuator_defaults(
     num_envs: int,
     num_joints: int,
     dof_offset: int,
+    env_stride: int,
     device: str,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | slice]:
     """Snapshot the initial kp/kd of every Newton actuator owned by one articulation.
@@ -255,6 +256,12 @@ def build_newton_actuator_defaults(
         num_joints: Articulation-local joint count.
         dof_offset: Offset of this articulation's DOFs in the env-major
             global index space (``0`` on PhysX, view-dependent on Newton).
+        env_stride: Whole-model per-env DOF count — the stride used to build
+            each actuator's env-major ``indices``. Equals ``num_joints`` on
+            PhysX, but exceeds it by the free-root DOFs on a floating-base
+            Newton articulation, so it must be passed explicitly rather than
+            assumed equal to ``num_joints``. The owning adapter's
+            :attr:`NewtonActuatorAdapter.num_joints` is exactly this value.
         device: Warp device string (e.g. ``"cuda:0"``).
 
     Returns:
@@ -290,14 +297,14 @@ def build_newton_actuator_defaults(
             wp.launch(
                 scatter_gain_kernel,
                 dim=act.indices.shape[0],
-                inputs=[ctrl.kp, flat_stiffness, act.indices, dof_offset, num_joints],
+                inputs=[ctrl.kp, flat_stiffness, act.indices, dof_offset, num_joints, env_stride],
                 device=wp_device,
             )
         if hasattr(ctrl, "kd"):
             wp.launch(
                 scatter_gain_kernel,
                 dim=act.indices.shape[0],
-                inputs=[ctrl.kd, flat_damping, act.indices, dof_offset, num_joints],
+                inputs=[ctrl.kd, flat_damping, act.indices, dof_offset, num_joints, env_stride],
                 device=wp_device,
             )
     stiffness = wp.to_torch(flat_stiffness.reshape((num_envs, num_joints)))

source/isaaclab_newton/isaaclab_newton/actuators/kernels.py

@@ -63,6 +63,7 @@ def scatter_gain_kernel(
     indices: wp.array(dtype=wp.uint32),
     dof_offset: int,
     num_joints: int,
+    env_stride: int,
 ):
     """Scatter per-actuator ``src`` values into a flat per-env-per-DOF ``dst``.
 
@@ -71,11 +72,30 @@ def scatter_gain_kernel(
     that ``randomize_actuator_gains`` reads as
     ``actuator.stiffness`` / ``.damping`` for its
     ``default_joint_stiffness`` / ``default_joint_damping`` baseline.
+
+    The actuator's ``indices`` are global DOF ids laid out env-major with a
+    per-env stride of ``env_stride`` — the *whole model's* per-env DOF count,
+    which on a floating-base articulation exceeds ``num_joints`` (the
+    articulation-local, actuated joint count) by the free-root DOFs. The env
+    index must therefore be decoded with ``env_stride``, not ``num_joints``;
+    the articulation-local joint offset is what remains after removing the
+    env's block and lands in ``[0, num_joints)`` because ``indices`` only ever
+    holds this articulation's joints.
+
+    Args:
+        src: Per-actuator parameter values (e.g. ``controller.kp``).
+        dst: Flat ``(num_envs * num_joints)`` articulation-local snapshot buffer.
+        indices: Actuator's flat env-major global DOF indices.
+        dof_offset: Offset of this articulation's DOFs in the env-major
+            global index space (``0`` on PhysX, view-dependent on Newton).
+        num_joints: Articulation-local joint count (``dst``'s inner stride).
+        env_stride: Whole-model per-env DOF count (the stride used to build
+            ``indices``).
     """
     i = wp.tid()
     global_dof = int(indices[i]) - dof_offset
-    env = global_dof // num_joints
-    local_dof = global_dof % num_joints
+    env = global_dof // env_stride
+    local_dof = global_dof - env * env_stride
     dst[env * num_joints + local_dof] = src[i]
 
 
@@ -97,6 +117,16 @@ def patch_actuator_param_kernel(
     ``dst[i]`` (the controller parameter) with ``values[e_pos, j_pos]``.
     Cells outside the sub-grid are left untouched.
 
+    Note:
+        This kernel is PhysX-only (the Newton backend patches gains via
+        :meth:`ArticulationView.set_actuator_parameter`). On PhysX every
+        joint's coordinate count equals its DOF count, so the per-env stride
+        used to build ``indices`` equals ``num_joints`` and the ``env`` /
+        ``joint`` split below is exact. Do not reuse this kernel on a layout
+        whose per-env DOF stride exceeds ``num_joints`` (e.g. a floating-base
+        Newton model) without threading the true stride, or the ``joint``
+        split will alias across envs — see :func:`scatter_gain_kernel`.
+
     Args:
         indices: Actuator's flat indices into the (env-major) DOF layout.
         env_id_pos: ``env_id_pos[env]`` gives the row in ``values`` for

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

@@ -3745,6 +3745,7 @@ def _process_actuators_cfg(self):
                     num_envs=self.num_instances,
                     num_joints=self.num_joints,
                     dof_offset=arti_start,
+                    env_stride=adapter.num_joints,
                     device=self.device,
                 )
             else:

source/isaaclab_newton/test/assets/test_newton_actuators_newton.py

@@ -693,6 +693,64 @@ def test_two_articulations(self):
                 torch.testing.assert_close(cp_kd_after[env_idx], cp_kd_before[env_idx])
 
 
+class TestNewtonActuatorGainSnapshotEnvStride(unittest.TestCase):
+    """Regression: the init-time kp/kd snapshot must be correct for every env.
+
+    ``build_newton_actuator_defaults`` scatters each Newton actuator's
+    ``controller.kp`` / ``controller.kd`` into a per-articulation
+    ``(num_envs, num_joints)`` tensor (``newton_default_stiffness`` /
+    ``newton_default_damping``), which ``randomize_actuator_gains`` reads as
+    its DR baseline. On a floating-base articulation the actuator ``indices``
+    are laid out env-major with a per-env stride equal to the *whole model's*
+    per-env DOF count (free-root DOFs + joints), which exceeds
+    ``articulation.num_joints``. If the scatter decodes the env with
+    ``num_joints`` instead of that stride, env 1's DOFs alias to the wrong
+    rows (and partly out of bounds), corrupting the snapshot for every env
+    past the first.
+
+    ANYmal-C is floating base (6 free-root DOFs + 12 actuated joints -> a
+    per-env stride of 18 vs. ``num_joints == 12``), so the bug manifests here
+    with ``NUM_ENVS == 2``: without the fix, ``newton_default_stiffness[1]``
+    is not uniformly the configured gain (its leading entries stay zero, as
+    they are never written).
+    """
+
+    def test_snapshot_matches_config_for_all_envs(self):
+        sim_cfg = SimulationCfg(dt=DT, physics=NEWTON_CFG, use_newton_actuators=True)
+        with build_simulation_context(
+            device="cuda:0",
+            gravity_enabled=True,
+            add_ground_plane=True,
+            sim_cfg=sim_cfg,
+        ) as sim:
+            sim._app_control_on_stop_handle = None
+            for i in range(NUM_ENVS):
+                sim_utils.create_prim(f"/World/Env_{i}", "Xform", translation=(i * 3.0, 0, 0))
+            art_cfg = ANYMAL_C_CFG.replace(
+                actuators=IDEAL_PD_ACTUATORS,
+                prim_path="/World/Env_.*/Robot",
+            )
+            anymal = Articulation(art_cfg)
+            sim.reset()
+            assert anymal.is_initialized
+
+            stiffness = anymal.newton_default_stiffness
+            damping = anymal.newton_default_damping
+            self.assertIsNotNone(stiffness, "expected a Newton kp snapshot with use_newton_actuators=True")
+            self.assertIsNotNone(damping, "expected a Newton kd snapshot with use_newton_actuators=True")
+
+            n_j = anymal.num_joints
+            self.assertEqual(tuple(stiffness.shape), (NUM_ENVS, n_j))
+            self.assertEqual(tuple(damping.shape), (NUM_ENVS, n_j))
+
+            # IDEAL_PD_ACTUATORS covers all 12 joints with constant gains, so
+            # every cell of both env rows must equal the configured value.
+            expected_kp = torch.full((NUM_ENVS, n_j), 40.0, device=anymal.device)
+            expected_kd = torch.full((NUM_ENVS, n_j), 5.0, device=anymal.device)
+            torch.testing.assert_close(stiffness, expected_kp)
+            torch.testing.assert_close(damping, expected_kd)
+
+
 # ---------------------------------------------------------------------------
 # DelayedPD equivalence: PD with actuator command delay
 # ---------------------------------------------------------------------------

source/isaaclab_physx/changelog.d/fix-actuator-gain-env-stride.skip

@@ -0,0 +1,4 @@
+Adapts the PhysX Newton-actuator call site to the new required ``env_stride``
+argument of ``build_newton_actuator_defaults``. No user-facing behavior change:
+PhysX articulations have no free or ball joints, so the per-env DOF stride
+already equals the articulation-local joint count.

source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py

@@ -4080,6 +4080,7 @@ def _process_actuators_cfg(self):
                     num_envs=self.num_instances,
                     num_joints=self.num_joints,
                     dof_offset=0,
+                    env_stride=adapter.num_joints,
                     device=self.device,
                 )
                 self.write_joint_stiffness_to_sim_index(stiffness=0.0, joint_ids=adapter.joint_indices)