From cd455c7eec615882451f38a2ab38b54206276d9c Mon Sep 17 00:00:00 2001 From: Antoine Richard Date: Wed, 8 Jul 2026 09:11:22 +0200 Subject: [PATCH] Fix actuator gain snapshot env stride decode The init-time gain snapshot kernel decoded environment and local joint slots from global model DOF indices using the articulation's joint count as the per-environment stride. Newton's global DOF space also contains free-root DOFs, so for floating-base articulations with more than one environment the decode misattributed later environments' joints and wrote past the snapshot buffer: silently corrupted gain randomization baselines on CUDA, and heap corruption aborts on CPU under newer Warp allocator layouts. Thread the true per-environment DOF stride from the adapter into the kernel; the PhysX call site passes an identical stride and is behaviorally unchanged. The existing suites missed this because the equivalence tests never read the snapshot and the randomization test reads controller parameters directly. The new regression test asserts the snapshot values for every environment row of a floating-base multi-environment articulation and fails deterministically without the fix. --- .../fix-actuator-gain-env-stride.rst | 12 ++++ .../isaaclab_newton/actuators/adapter.py | 11 +++- .../isaaclab_newton/actuators/kernels.py | 34 ++++++++++- .../assets/articulation/articulation.py | 1 + .../assets/test_newton_actuators_newton.py | 58 +++++++++++++++++++ .../fix-actuator-gain-env-stride.skip | 4 ++ .../assets/articulation/articulation.py | 1 + 7 files changed, 117 insertions(+), 4 deletions(-) create mode 100644 source/isaaclab_newton/changelog.d/fix-actuator-gain-env-stride.rst create mode 100644 source/isaaclab_physx/changelog.d/fix-actuator-gain-env-stride.skip diff --git a/source/isaaclab_newton/changelog.d/fix-actuator-gain-env-stride.rst b/source/isaaclab_newton/changelog.d/fix-actuator-gain-env-stride.rst new file mode 100644 index 000000000000..164036dccf92 --- /dev/null +++ b/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. diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py b/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py index 0d28cf576dfb..da1de70d408d 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/adapter.py +++ b/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))) diff --git a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py b/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py index 8205773b640a..7c9312ede2a6 100644 --- a/source/isaaclab_newton/isaaclab_newton/actuators/kernels.py +++ b/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 diff --git a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py b/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py index 2fb62bb9a29a..8d7d9e5ef4a3 100644 --- a/source/isaaclab_newton/isaaclab_newton/assets/articulation/articulation.py +++ b/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: diff --git a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py b/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py index 68531c270481..f5457690f9e2 100644 --- a/source/isaaclab_newton/test/assets/test_newton_actuators_newton.py +++ b/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 # --------------------------------------------------------------------------- diff --git a/source/isaaclab_physx/changelog.d/fix-actuator-gain-env-stride.skip b/source/isaaclab_physx/changelog.d/fix-actuator-gain-env-stride.skip new file mode 100644 index 000000000000..2680d9068d7f --- /dev/null +++ b/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. diff --git a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py b/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py index 39d9604e3922..2e6ebc16f80a 100644 --- a/source/isaaclab_physx/isaaclab_physx/assets/articulation/articulation.py +++ b/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)