Adds position to set external forces and torques (#1680)

# Description

This PR adds the option to set wrench positions to
`set_external_force_and_torque`. This is a non-breaking change as the
positions are passed as an optional argument. When no positions are set,
the function defaults to the original implementation, that is, no
positions are passed to PhysX. The PR also adds tests to check that the
position values are correctly set into their buffer, but does not check
if the resulting wrenches are correct. I did test the Quadcopter task
before and after this PR and the training results are exactly the same.

As of now, the function follows the original layout. But it could make
sense to offer the option to set the position in either the link frame
or the CoM frame. This would follow the recent changes made to the
set_pose and set_velocity methods for instance. However, this would be a
breaking change. Hence, for now, this has not been implemented. One
could also argue, that this could be done prior to feeding the positions
outside this method. Please let me know what you feel is best, and I'll
update the PR accordingly.

If one wanted to test the resulting wrenches, it would require a simple
test articulation like a 1kg sphere that would be used to accurately
compute the expected velocities. This is also feasible, but I feel like
this test is more on the PhysX side of things, let me know.

This change will require an update of the API documentation to include
the position argument.

Fixes #1678 

## Type of change

- New feature (non-breaking change which adds functionality)
- This change requires a documentation update

## Checklist

- [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
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

---------

Co-authored-by: Kelly Guo <[email protected]>

Author: AntoineRichard

source/isaaclab/isaaclab/assets/articulation/articulation.py

@@ -169,6 +169,7 @@ def reset(self, env_ids: Sequence[int] | None = None):
         # reset external wrench
         self._external_force_b[env_ids] = 0.0
         self._external_torque_b[env_ids] = 0.0
+        self._external_wrench_positions_b[env_ids] = 0.0
 
     def write_data_to_sim(self):
         """Write external wrenches and joint commands to the simulation.
@@ -182,13 +183,22 @@ def write_data_to_sim(self):
         """
         # write external wrench
         if self.has_external_wrench:
-            self.root_physx_view.apply_forces_and_torques_at_position(
-                force_data=self._external_force_b.view(-1, 3),
-                torque_data=self._external_torque_b.view(-1, 3),
-                position_data=None,
-                indices=self._ALL_INDICES,
-                is_global=False,
-            )
+            if self.uses_external_wrench_positions:
+                self.root_physx_view.apply_forces_and_torques_at_position(
+                    force_data=self._external_force_b.view(-1, 3),
+                    torque_data=self._external_torque_b.view(-1, 3),
+                    position_data=self._external_wrench_positions_b.view(-1, 3),
+                    indices=self._ALL_INDICES,
+                    is_global=self._use_global_wrench_frame,
+                )
+            else:
+                self.root_physx_view.apply_forces_and_torques_at_position(
+                    force_data=self._external_force_b.view(-1, 3),
+                    torque_data=self._external_torque_b.view(-1, 3),
+                    position_data=None,
+                    indices=self._ALL_INDICES,
+                    is_global=self._use_global_wrench_frame,
+                )
 
         # apply actuator models
         self._apply_actuator_model()
@@ -829,14 +839,16 @@ def set_external_force_and_torque(
         self,
         forces: torch.Tensor,
         torques: torch.Tensor,
+        positions: torch.Tensor | None = None,
         body_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Set external force and torque to apply on the asset's bodies in their local frame.
 
         For many applications, we want to keep the applied external force on rigid bodies constant over a period of
         time (for instance, during the policy control). This function allows us to store the external force and torque
-        into buffers which are then applied to the simulation at every step.
+        into buffers which are then applied to the simulation at every step. Optionally, set the position to apply the
+        external wrench at (in the local link frame of the bodies).
 
         .. caution::
             If the function is called with empty forces and torques, then this function disables the application
@@ -855,6 +867,7 @@ def set_external_force_and_torque(
         Args:
             forces: External forces in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3).
             torques: External torques in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3).
+            positions: Positions to apply external wrench. Shape is (len(env_ids), len(body_ids), 3). Defaults to None.
             body_ids: Body indices to apply external wrench to. Defaults to None (all bodies).
             env_ids: Environment indices to apply external wrench to. Defaults to None (all instances).
         """
@@ -887,6 +900,17 @@ def set_external_force_and_torque(
         self._external_force_b.flatten(0, 1)[indices] = forces.flatten(0, 1)
         self._external_torque_b.flatten(0, 1)[indices] = torques.flatten(0, 1)
 
+        # If the positions are not provided, the behavior and performance of the simulation should not be affected.
+        if positions is not None:
+            # Generates a flag that is set for a full simulation step. This is done to avoid discarding
+            # the external wrench positions when multiple calls to this functions are made with and without positions.
+            self.uses_external_wrench_positions = True
+            self._external_wrench_positions_b.flatten(0, 1)[indices] = positions.flatten(0, 1)
+        else:
+            # If the positions are not provided, and the flag is set, then we need to ensure that the desired positions are zeroed.
+            if self.uses_external_wrench_positions:
+                self._external_wrench_positions_b.flatten(0, 1)[indices] = 0.0
+
     def set_joint_position_target(
         self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None
     ):
@@ -1229,8 +1253,10 @@ def _create_buffers(self):
 
         # external forces and torques
         self.has_external_wrench = False
+        self.uses_external_wrench_positions = False
         self._external_force_b = torch.zeros((self.num_instances, self.num_bodies, 3), device=self.device)
         self._external_torque_b = torch.zeros_like(self._external_force_b)
+        self._external_wrench_positions_b = torch.zeros_like(self._external_force_b)
 
         # asset named data
         self._data.joint_names = self.joint_names
@@ -1298,6 +1324,15 @@ def _process_cfg(self):
         default_root_state = torch.tensor(default_root_state, dtype=torch.float, device=self.device)
         self._data.default_root_state = default_root_state.repeat(self.num_instances, 1)
 
+        # -- external wrench
+        external_wrench_frame = self.cfg.articulation_external_wrench_frame
+        if external_wrench_frame == "local":
+            self._use_global_wrench_frame = False
+        elif external_wrench_frame == "world":
+            self._use_global_wrench_frame = True
+        else:
+            raise ValueError(f"Invalid external wrench frame: {external_wrench_frame}. Must be 'local' or 'world'.")
+
         # -- joint state
         self._data.default_joint_pos = torch.zeros(self.num_instances, self.num_joints, device=self.device)
         self._data.default_joint_vel = torch.zeros_like(self._data.default_joint_pos)

source/isaaclab/isaaclab/assets/articulation/articulation_cfg.py

@@ -44,6 +44,13 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg):
     If not provided will search for a prim with the ArticulationRootAPI. Should start with a slash.
     """
 
+    articulation_external_wrench_frame: str = "local"
+    """Frame in which external wrenches are applied. Defaults to "local".
+
+    If "local", the external wrenches are applied in the local frame of the articulation root.
+    If "world", the external wrenches are applied in the world frame.
+    """
+
     init_state: InitialStateCfg = InitialStateCfg()
     """Initial state of the articulated object. Defaults to identity pose with zero velocity and zero joint state."""
 

source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py

@@ -103,6 +103,7 @@ def reset(self, env_ids: Sequence[int] | None = None):
         # reset external wrench
         self._external_force_b[env_ids] = 0.0
         self._external_torque_b[env_ids] = 0.0
+        self._external_wrench_positions_b[env_ids] = 0.0
 
     def write_data_to_sim(self):
         """Write external wrench to the simulation.
@@ -113,13 +114,22 @@ def write_data_to_sim(self):
         """
         # write external wrench
         if self.has_external_wrench:
-            self.root_physx_view.apply_forces_and_torques_at_position(
-                force_data=self._external_force_b.view(-1, 3),
-                torque_data=self._external_torque_b.view(-1, 3),
-                position_data=None,
-                indices=self._ALL_INDICES,
-                is_global=False,
-            )
+            if self.uses_external_wrench_positions:
+                self.root_physx_view.apply_forces_and_torques_at_position(
+                    force_data=self._external_force_b.view(-1, 3),
+                    torque_data=self._external_torque_b.view(-1, 3),
+                    position_data=self._external_wrench_positions_b.view(-1, 3),
+                    indices=self._ALL_INDICES,
+                    is_global=self._use_global_wrench_frame,
+                )
+            else:
+                self.root_physx_view.apply_forces_and_torques_at_position(
+                    force_data=self._external_force_b.view(-1, 3),
+                    torque_data=self._external_torque_b.view(-1, 3),
+                    position_data=None,
+                    indices=self._ALL_INDICES,
+                    is_global=self._use_global_wrench_frame,
+                )
 
     def update(self, dt: float):
         self._data.update(dt)
@@ -357,14 +367,16 @@ def set_external_force_and_torque(
         self,
         forces: torch.Tensor,
         torques: torch.Tensor,
+        positions: torch.Tensor | None = None,
         body_ids: Sequence[int] | slice | None = None,
         env_ids: Sequence[int] | None = None,
     ):
         """Set external force and torque to apply on the asset's bodies in their local frame.
 
         For many applications, we want to keep the applied external force on rigid bodies constant over a period of
         time (for instance, during the policy control). This function allows us to store the external force and torque
-        into buffers which are then applied to the simulation at every step.
+        into buffers which are then applied to the simulation at every step. Optionally, set the position to apply the
+        external wrench at (in the local link frame of the bodies).
 
         .. caution::
             If the function is called with empty forces and torques, then this function disables the application
@@ -383,6 +395,7 @@ def set_external_force_and_torque(
         Args:
             forces: External forces in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3).
             torques: External torques in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3).
+            positions: External wrench positions in bodies' local frame. Shape is (len(env_ids), len(body_ids), 3). Defaults to None.
             body_ids: Body indices to apply external wrench to. Defaults to None (all bodies).
             env_ids: Environment indices to apply external wrench to. Defaults to None (all instances).
         """
@@ -407,6 +420,13 @@ def set_external_force_and_torque(
         self._external_force_b[env_ids, body_ids] = forces
         self._external_torque_b[env_ids, body_ids] = torques
 
+        if positions is not None:
+            self.uses_external_wrench_positions = True
+            self._external_wrench_positions_b[env_ids, body_ids] = positions
+        else:
+            if self.uses_external_wrench_positions:
+                self._external_wrench_positions_b[env_ids, body_ids] = 0.0
+
     """
     Internal helper.
     """
@@ -483,6 +503,8 @@ def _create_buffers(self):
         self.has_external_wrench = False
         self._external_force_b = torch.zeros((self.num_instances, self.num_bodies, 3), device=self.device)
         self._external_torque_b = torch.zeros_like(self._external_force_b)
+        self.uses_external_wrench_positions = False
+        self._external_wrench_positions_b = torch.zeros_like(self._external_force_b)
 
         # set information about rigid body into data
         self._data.body_names = self.body_names
@@ -503,6 +525,15 @@ def _process_cfg(self):
         default_root_state = torch.tensor(default_root_state, dtype=torch.float, device=self.device)
         self._data.default_root_state = default_root_state.repeat(self.num_instances, 1)
 
+        # -- external wrench
+        external_wrench_frame = self.cfg.object_external_wrench_frame
+        if external_wrench_frame == "local":
+            self._use_global_wrench_frame = False
+        elif external_wrench_frame == "world":
+            self._use_global_wrench_frame = True
+        else:
+            raise ValueError(f"Invalid external wrench frame: {external_wrench_frame}. Must be 'local' or 'world'.")
+
     """
     Internal simulation callbacks.
     """

source/isaaclab/isaaclab/assets/rigid_object/rigid_object_cfg.py

@@ -30,3 +30,10 @@ class InitialStateCfg(AssetBaseCfg.InitialStateCfg):
 
     init_state: InitialStateCfg = InitialStateCfg()
     """Initial state of the rigid object. Defaults to identity pose with zero velocity."""
+
+    object_external_wrench_frame: str = "local"
+    """Frame in which external wrenches are applied. Defaults to "local".
+
+    If "local", the external wrenches are applied in the local frame of the articulation root.
+    If "world", the external wrenches are applied in the world frame.
+    """

source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection.py

@@ -165,6 +165,7 @@ def reset(self, env_ids: torch.Tensor | None = None, object_ids: slice | torch.T
         # reset external wrench
         self._external_force_b[env_ids[:, None], object_ids] = 0.0
         self._external_torque_b[env_ids[:, None], object_ids] = 0.0
+        self._external_wrench_positions_b[env_ids[:, None], object_ids] = 0.0
 
     def write_data_to_sim(self):
         """Write external wrench to the simulation.
@@ -175,13 +176,22 @@ def write_data_to_sim(self):
         """
         # write external wrench
         if self.has_external_wrench:
-            self.root_physx_view.apply_forces_and_torques_at_position(
-                force_data=self.reshape_data_to_view(self._external_force_b),
-                torque_data=self.reshape_data_to_view(self._external_torque_b),
-                position_data=None,
-                indices=self._env_obj_ids_to_view_ids(self._ALL_ENV_INDICES, self._ALL_OBJ_INDICES),
-                is_global=False,
-            )
+            if self.uses_external_wrench_positions:
+                self.root_physx_view.apply_forces_and_torques_at_position(
+                    force_data=self.reshape_data_to_view(self._external_force_b),
+                    torque_data=self.reshape_data_to_view(self._external_torque_b),
+                    position_data=self.reshape_data_to_view(self._external_wrench_positions_b),
+                    indices=self._env_obj_ids_to_view_ids(self._ALL_ENV_INDICES, self._ALL_OBJ_INDICES),
+                    is_global=self._use_global_wrench_frame,
+                )
+            else:
+                self.root_physx_view.apply_forces_and_torques_at_position(
+                    force_data=self.reshape_data_to_view(self._external_force_b),
+                    torque_data=self.reshape_data_to_view(self._external_torque_b),
+                    position_data=None,
+                    indices=self._env_obj_ids_to_view_ids(self._ALL_ENV_INDICES, self._ALL_OBJ_INDICES),
+                    is_global=self._use_global_wrench_frame,
+                )
 
     def update(self, dt: float):
         self._data.update(dt)
@@ -486,6 +496,7 @@ def set_external_force_and_torque(
         self,
         forces: torch.Tensor,
         torques: torch.Tensor,
+        positions: torch.Tensor | None = None,
         object_ids: slice | torch.Tensor | None = None,
         env_ids: torch.Tensor | None = None,
     ):
@@ -512,6 +523,7 @@ def set_external_force_and_torque(
         Args:
             forces: External forces in bodies' local frame. Shape is (len(env_ids), len(object_ids), 3).
             torques: External torques in bodies' local frame. Shape is (len(env_ids), len(object_ids), 3).
+            positions: External wrench positions in bodies' local frame. Shape is (len(env_ids), len(object_ids), 3).
             object_ids: Object indices to apply external wrench to. Defaults to None (all objects).
             env_ids: Environment indices to apply external wrench to. Defaults to None (all instances).
         """
@@ -532,6 +544,12 @@ def set_external_force_and_torque(
         # set into internal buffers
         self._external_force_b[env_ids[:, None], object_ids] = forces
         self._external_torque_b[env_ids[:, None], object_ids] = torques
+        if positions is not None:
+            self.uses_external_wrench_positions = True
+            self._external_wrench_positions_b[env_ids[:, None], object_ids] = positions
+        else:
+            if self.uses_external_wrench_positions:
+                self._external_wrench_positions_b[env_ids[:, None], object_ids] = 0.0
 
     """
     Helper functions.
@@ -643,6 +661,8 @@ def _create_buffers(self):
         self.has_external_wrench = False
         self._external_force_b = torch.zeros((self.num_instances, self.num_objects, 3), device=self.device)
         self._external_torque_b = torch.zeros_like(self._external_force_b)
+        self._external_wrench_positions_b = torch.zeros_like(self._external_force_b)
+        self.uses_external_wrench_positions = False
 
         # set information about rigid body into data
         self._data.object_names = self.object_names
@@ -671,6 +691,15 @@ def _process_cfg(self):
         default_object_states = torch.cat(default_object_states, dim=1)
         self._data.default_object_state = default_object_states
 
+        # -- external wrench
+        external_wrench_frame = self.cfg.objects_external_wrench_frame
+        if external_wrench_frame == "local":
+            self._use_global_wrench_frame = False
+        elif external_wrench_frame == "world":
+            self._use_global_wrench_frame = True
+        else:
+            raise ValueError(f"Invalid external wrench frame: {external_wrench_frame}. Must be 'local' or 'world'.")
+
     def _env_obj_ids_to_view_ids(
         self, env_ids: torch.Tensor, object_ids: Sequence[int] | slice | torch.Tensor
     ) -> torch.Tensor:

source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection_cfg.py

@@ -26,3 +26,10 @@ class RigidObjectCollectionCfg:
 
     The keys are the names for the objects, which are used as unique identifiers throughout the code.
     """
+
+    objects_external_wrench_frame: str = "local"
+    """Frame in which external wrenches are applied. Defaults to "local".
+
+    If "local", the external wrenches are applied in the local frame of the articulation root.
+    If "world", the external wrenches are applied in the world frame.
+    """