Implement randomize_actuator_stiffness_and_dampning for explicit actuator models. #1000

Set default stiffness, damping for explicit actuators #1006

fixes randomize_joint_parameters shape mismatch #1011

Author: MoreTore

CONTRIBUTORS.md

@@ -74,7 +74,7 @@ Guidelines for modifications:
 * Zhengyu Zhang
 * Ziqi Fan
 * Qian Wan
-* Wei Yang
+* Ryley McCarroll
 
 ## Acknowledgements
 

source/extensions/omni.isaac.lab/omni/isaac/lab/actuators/actuator_base.py

@@ -207,6 +207,7 @@ def _parse_joint_parameter(
             TypeError: If the parameter value is not of the expected type.
             TypeError: If the default value is not of the expected type.
             ValueError: If the parameter value is None and no default value is provided.
+            ValueError: If the default value tensor is the wrong shape.
         """
         # create parameter buffer
         param = torch.zeros(self._num_envs, self.num_joints, device=self._device)
@@ -231,7 +232,14 @@ def _parse_joint_parameter(
                 param[:] = float(default_value)
             elif isinstance(default_value, torch.Tensor):
                 # if tensor, then use the same tensor for all joints
-                param[:] = default_value.float()
+                if default_value.shape == (self._num_envs, self.num_joints):
+                    param = default_value.float()
+                else:
+                    raise ValueError(
+                        "Invalid default value tensor shape.\n"
+                        f"Got: {default_value.shape}\n"
+                        f"Expected: {(self._num_envs, self.num_joints)}"
+                    )
             else:
                 raise TypeError(
                     f"Invalid type for default value: {type(default_value)} for "

source/extensions/omni.isaac.lab/omni/isaac/lab/assets/articulation/articulation.py

@@ -1105,12 +1105,10 @@ def _process_actuators_cfg(self):
         self._has_implicit_actuators = False
 
         # cache the values coming from the usd
-        usd_stiffness = self.root_physx_view.get_dof_stiffnesses().clone()
-        usd_damping = self.root_physx_view.get_dof_dampings().clone()
-        usd_armature = self.root_physx_view.get_dof_armatures().clone()
-        usd_friction = self.root_physx_view.get_dof_friction_coefficients().clone()
-        usd_effort_limit = self.root_physx_view.get_dof_max_forces().clone()
-        usd_velocity_limit = self.root_physx_view.get_dof_max_velocities().clone()
+        self._data.default_joint_stiffness = self.root_physx_view.get_dof_stiffnesses().to(self.device).clone()
+        self._data.default_joint_damping = self.root_physx_view.get_dof_dampings().to(self.device).clone()
+        self._data.default_joint_armature = self.root_physx_view.get_dof_armatures().to(self.device).clone()
+        self._data.default_joint_friction = self.root_physx_view.get_dof_friction_coefficients().to(self.device).clone()
 
         # iterate over all actuator configurations
         for actuator_name, actuator_cfg in self.cfg.actuators.items():
@@ -1134,12 +1132,12 @@ def _process_actuators_cfg(self):
                 ),
                 num_envs=self.num_instances,
                 device=self.device,
-                stiffness=usd_stiffness[:, joint_ids],
-                damping=usd_damping[:, joint_ids],
-                armature=usd_armature[:, joint_ids],
-                friction=usd_friction[:, joint_ids],
-                effort_limit=usd_effort_limit[:, joint_ids],
-                velocity_limit=usd_velocity_limit[:, joint_ids],
+                stiffness=self._data.default_joint_stiffness[:, joint_ids],
+                damping=self._data.default_joint_damping[:, joint_ids],
+                armature=self._data.default_joint_armature[:, joint_ids],
+                friction=self._data.default_joint_friction[:, joint_ids],
+                effort_limit=self.root_physx_view.get_dof_max_forces().to(self.device).clone()[:, joint_ids],
+                velocity_limit=self.root_physx_view.get_dof_max_velocities().to(self.device).clone()[:, joint_ids],
             )
             # log information on actuator groups
             omni.log.info(
@@ -1165,14 +1163,9 @@ def _process_actuators_cfg(self):
                 self.write_joint_effort_limit_to_sim(1.0e9, joint_ids=actuator.joint_indices)
                 self.write_joint_armature_to_sim(actuator.armature, joint_ids=actuator.joint_indices)
                 self.write_joint_friction_to_sim(actuator.friction, joint_ids=actuator.joint_indices)
-
-        # set the default joint parameters based on the changes from the actuators
-        self._data.default_joint_stiffness = self.root_physx_view.get_dof_stiffnesses().to(device=self.device).clone()
-        self._data.default_joint_damping = self.root_physx_view.get_dof_dampings().to(device=self.device).clone()
-        self._data.default_joint_armature = self.root_physx_view.get_dof_armatures().to(device=self.device).clone()
-        self._data.default_joint_friction = (
-            self.root_physx_view.get_dof_friction_coefficients().to(device=self.device).clone()
-        )
+                # Store the actual default stiffness and damping values for explicit actuators (not written the sim)
+                self._data.default_joint_stiffness[:, actuator.joint_indices] = actuator.stiffness
+                self._data.default_joint_damping[:, actuator.joint_indices] = actuator.damping
 
         # perform some sanity checks to ensure actuators are prepared correctly
         total_act_joints = sum(actuator.num_joints for actuator in self.actuators.values())

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py

@@ -300,54 +300,53 @@ def randomize_actuator_gains(
     .. tip::
         For implicit actuators, this function uses CPU tensors to assign the actuator gains into the simulation.
         In such cases, it is recommended to use this function only during the initialization of the environment.
-
-    Raises:
-        NotImplementedError: If the joint indices are in explicit motor mode. This operation is currently
-            not supported for explicit actuator models.
     """
-    # extract the used quantities (to enable type-hinting)
+    # Extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
 
-    # resolve environment ids
+    # Resolve environment ids
     if env_ids is None:
         env_ids = torch.arange(env.scene.num_envs, device=asset.device)
 
-    # resolve joint indices
-    if asset_cfg.joint_ids == slice(None):
-        joint_ids_list = range(asset.num_joints)
-        joint_ids = slice(None)  # for optimization purposes
-    else:
-        joint_ids_list = asset_cfg.joint_ids
-        joint_ids = torch.tensor(asset_cfg.joint_ids, dtype=torch.int, device=asset.device)
+    joint_ids_list = list(range(asset.num_joints)) if isinstance(asset_cfg.joint_ids, slice) else asset_cfg.joint_ids
 
-    # check if none of the joint indices are in explicit motor mode
-    for joint_index in joint_ids_list:
-        for act_name, actuator in asset.actuators.items():
-            # if joint indices are a slice (i.e., all joints are captured) or the joint index is in the actuator
-            if actuator.joint_indices == slice(None) or joint_index in actuator.joint_indices:
-                if not isinstance(actuator, ImplicitActuator):
-                    raise NotImplementedError(
-                        "Event term 'randomize_actuator_stiffness_and_damping' is performed on asset"
-                        f" '{asset_cfg.name}' on the joint '{asset.joint_names[joint_index]}' ('{joint_index}') which"
-                        f" uses an explicit actuator model '{act_name}<{actuator.__class__.__name__}>'. This operation"
-                        " is currently not supported for explicit actuator models."
-                    )
+    def randomize(data: torch.Tensor, params: tuple[float, float]) -> torch.Tensor:
+        return _randomize_prop_by_op(
+            data, params, env_ids, target_joint_indices, operation=operation, distribution=distribution
+        )
 
-    # sample joint properties from the given ranges and set into the physics simulation
-    # -- stiffness
-    if stiffness_distribution_params is not None:
-        stiffness = asset.data.default_joint_stiffness.to(asset.device).clone()
-        stiffness = _randomize_prop_by_op(
-            stiffness, stiffness_distribution_params, env_ids, joint_ids, operation=operation, distribution=distribution
-        )[env_ids][:, joint_ids]
-        asset.write_joint_stiffness_to_sim(stiffness, joint_ids=joint_ids, env_ids=env_ids)
-    # -- damping
-    if damping_distribution_params is not None:
-        damping = asset.data.default_joint_damping.to(asset.device).clone()
-        damping = _randomize_prop_by_op(
-            damping, damping_distribution_params, env_ids, joint_ids, operation=operation, distribution=distribution
-        )[env_ids][:, joint_ids]
-        asset.write_joint_damping_to_sim(damping, joint_ids=joint_ids, env_ids=env_ids)
+    # Loop through actuators and randomize gains
+    for actuator in asset.actuators.values():
+        # Resolve joint indices for each actuator
+        # Only target the joints specified by the asset_cfg.joint_ids (joint_ids_list)
+        if isinstance(actuator.joint_indices, slice):  # All joints of the articulation are this actuator type
+            target_joint_indices = torch.tensor(joint_ids_list, dtype=torch.int, device=asset.device)
+        else:
+            target_joint_indices = torch.tensor(
+                list(set(actuator.joint_indices).intersection(joint_ids_list)), dtype=torch.int, device=asset.device
+            )
+        if not target_joint_indices.nbytes:  # Skip when there are no joints
+            continue
+
+        all_envs = torch.arange(env.scene.num_envs, device=asset.device)
+        # Randomize stiffness
+        if stiffness_distribution_params is not None:
+            stiffness = asset.data.default_joint_stiffness.to(asset.device).clone()
+            randomize(stiffness, stiffness_distribution_params)
+            actuator.stiffness = stiffness[all_envs[:, None], actuator.joint_indices]
+            if isinstance(actuator, ImplicitActuator):
+                asset.write_joint_stiffness_to_sim(
+                    stiffness[env_ids][:, target_joint_indices], joint_ids=target_joint_indices, env_ids=env_ids
+                )
+        # Randomize damping
+        if damping_distribution_params is not None:
+            damping = asset.data.default_joint_damping.to(asset.device).clone()
+            randomize(damping, damping_distribution_params)
+            actuator.damping = damping[all_envs[:, None], actuator.joint_indices]
+            if isinstance(actuator, ImplicitActuator):
+                asset.write_joint_damping_to_sim(
+                    damping[env_ids][:, target_joint_indices], joint_ids=target_joint_indices, env_ids=env_ids
+                )
 
 
 def randomize_joint_parameters(
@@ -383,7 +382,7 @@ def randomize_joint_parameters(
 
     # resolve joint indices
     if asset_cfg.joint_ids == slice(None):
-        joint_ids = slice(None)  # for optimization purposes
+        joint_ids = torch.arange(asset.num_joints, dtype=torch.int, device=asset.device)
     else:
         joint_ids = torch.tensor(asset_cfg.joint_ids, dtype=torch.int, device=asset.device)