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)