Adds WrenchComposer to handle temporary/permanent wrenches (isaac-sim#3287)

# Description

Adds the ability to compose forces onto rigid bodies, rigid body
collections and articulations.
This should help implement drones, boats, and satellites into the
framework.

## Usage

```python

# Permanent forces can now be composed:
# Adding two forces in a single step on the same body
asset.permanent_wrench_composer.set_forces_and_torques(forces=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0])
# Compose local and global forces together
asset.permanent_wrench_composer.add_forces_and_torques(forces=torch.ones(1, 1, 3), env_ids=[0], object_ids=[1], is_global=True)
# Adding torques to the same body
asset.permanent_wrench_composer.add_forces_and_torques(torques=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0])
#Adding forces and torques to the same body
asset.permanent_wrench_composer.add_forces_and_torques(forces=torch.ones(1, 1, 3), torques=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0])
# Adding forces and torques to the same body with different positions
asset.permanent_wrench_composer.add_forces_and_torques(forces=torch.ones(1, 1, 3), torques=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0], positions=torch.ones(1, 1, 3))
# Adding forces and torques to the same body with different positions in the global frame. Note, it composes local and global wrenches seamlessly. 
asset.permanent_wrench_composer.add_forces_and_torques(forces=torch.ones(1, 1, 3), torques=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0], positions=torch.ones(1, 1, 3), is_global=True)


# We can now apply instantaneous wrenches that are only applied for a single simulation step:
asset.instantaneous_wrench_composer.add_forces_and_torques(forces=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0])
asset.instantaneous_wrench_composer.add_forces_and_torques(forces=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0])
asset.instantaneous_wrench_composer.add_forces_and_torques(forces=torch.ones(1, 1, 3), env_ids=[0], object_ids=[0])

# The instantaneous wrenches and the permanent wrenches are composed automatically when the wrenches are written to the simulation. The instantaneous wrenches are reseted after being written to the sim. 

```
Fixes isaac-sim#3286

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

---------

Signed-off-by: Antoine RICHARD <antoiner@nvidia.com>
Signed-off-by: Kelly Guo <kellyg@nvidia.com>
Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>

Author: 3 people

scripts/demos/quadcopter.py

@@ -101,7 +101,11 @@ def main():
         forces = torch.zeros(robot.num_instances, 4, 3, device=sim.device)
         torques = torch.zeros_like(forces)
         forces[..., 2] = robot_mass * gravity / 4.0
-        robot.set_external_force_and_torque(forces, torques, body_ids=prop_body_ids)
+        robot.permanent_wrench_composer.set_forces_and_torques(
+            forces=forces,
+            torques=torques,
+            body_ids=prop_body_ids,
+        )
         robot.write_data_to_sim()
         # perform step
         sim.step()

source/isaaclab/docs/CHANGELOG.rst

@@ -1,6 +1,30 @@
 Changelog
 ---------
 
+0.53.2 (2026-01-14)
+~~~~~~~~~~~~~~~~~~~
+
+Added
+^^^^^
+
+* Added :class:`~isaaclab.assets.utils.wrench_composer.WrenchComposer` to compose forces and torques at the body's center of mass frame.
+* Added :meth:`~isaaclab.assets.Articulation.instantaneous_wrench_composer` to add or set instantaneous external wrenches to the articulation.
+* Added :meth:`~isaaclab.assets.Articulation.permanent_wrench_composer` to add or set permanent external wrenches to the articulation.
+* Added :meth:`~isaaclab.assets.RigidObject.instantaneous_wrench_composer` to add or set instantaneous external wrenches to the rigid object.
+* Added :meth:`~isaaclab.assets.RigidObject.permanent_wrench_composer` to add or set permanent external wrenches to the rigid object.
+* Added :meth:`~isaaclab.assets.RigidObjectCollection.instantaneous_wrench_composer` to add or set instantaneous external wrenches to the rigid object collection.
+* Added :meth:`~isaaclab.assets.RigidObjectCollection.permanent_wrench_composer` to add or set permanent external wrenches to the rigid object collection.
+* Added unit tests for the wrench composer.
+* Added kernels for the wrench composer in the :mod:`isaaclab.utils.warp.kernels` module.
+
+Changed
+^^^^^^^
+
+* Deprecated :meth:`~isaaclab.assets.Articulation.set_external_force_and_torque`  in favor of :meth:`~isaaclab.assets.Articulation.permanent_wrench_composer.set_forces_and_torques`.
+* Deprecated :meth:`~isaaclab.assets.RigidObject.set_external_force_and_torque`  in favor of :meth:`~isaaclab.assets.RigidObject.permanent_wrench_composer.set_forces_and_torques`.
+* Deprecated :meth:`~isaaclab.assets.RigidObjectCollection.set_external_force_and_torque`  in favor of :meth:`~isaaclab.assets.RigidObjectCollection.permanent_wrench_composer.set_forces_and_torques`.
+* Modified the tests of the articulation, rigid object, and rigid object collection to use the new permanent and instantaneous external wrench functions and test them.
+
 0.53.1 (2026-01-08)
 ~~~~~~~~~~~~~~~~~~~
 

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

@@ -13,6 +13,7 @@
 from typing import TYPE_CHECKING
 
 import torch
+import warp as wp
 from prettytable import PrettyTable
 
 import omni.physics.tensors.impl.api as physx
@@ -25,6 +26,7 @@
 from isaaclab.actuators import ActuatorBase, ActuatorBaseCfg, ImplicitActuator
 from isaaclab.utils.types import ArticulationActions
 from isaaclab.utils.version import get_isaac_sim_version
+from isaaclab.utils.wrench_composer import WrenchComposer
 
 from ..asset_base import AssetBase
 from .articulation_data import ArticulationData
@@ -169,6 +171,35 @@ def root_physx_view(self) -> physx.ArticulationView:
         """
         return self._root_physx_view
 
+    @property
+    def instantaneous_wrench_composer(self) -> WrenchComposer:
+        """Instantaneous wrench composer.
+
+        Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench
+        composer are only valid for the current simulation step. At the end of the simulation step, the wrenches set
+        to this object are discarded. This is useful to apply forces that change all the time, things like drag forces
+        for instance.
+
+        Note:
+            Permanent wrenches are composed into the instantaneous wrench before the instantaneous wrenches are
+            applied to the simulation.
+        """
+        return self._instantaneous_wrench_composer
+
+    @property
+    def permanent_wrench_composer(self) -> WrenchComposer:
+        """Permanent wrench composer.
+
+        Returns a :class:`~isaaclab.utils.wrench_composer.WrenchComposer` instance. Wrenches added or set to this wrench
+        composer are persistent and are applied to the simulation at every step. This is useful to apply forces that
+        are constant over a period of time, things like the thrust of a motor for instance.
+
+        Note:
+            Permanent wrenches are composed into the instantaneous wrench before the instantaneous wrenches are
+            applied to the simulation.
+        """
+        return self._permanent_wrench_composer
+
     """
     Operations.
     """
@@ -180,10 +211,9 @@ def reset(self, env_ids: Sequence[int] | None = None):
         # reset actuators
         for actuator in self.actuators.values():
             actuator.reset(env_ids)
-        # 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
+        # reset external wrenches.
+        self._instantaneous_wrench_composer.reset(env_ids)
+        self._permanent_wrench_composer.reset(env_ids)
 
     def write_data_to_sim(self):
         """Write external wrenches and joint commands to the simulation.
@@ -196,23 +226,33 @@ def write_data_to_sim(self):
             This ensures that the external wrench is applied at every simulation step.
         """
         # write external wrench
-        if self.has_external_wrench:
-            if self.uses_external_wrench_positions:
+        if self._instantaneous_wrench_composer.active or self._permanent_wrench_composer.active:
+            if self._instantaneous_wrench_composer.active:
+                # Compose instantaneous wrench with permanent wrench
+                self._instantaneous_wrench_composer.add_forces_and_torques(
+                    forces=self._permanent_wrench_composer.composed_force,
+                    torques=self._permanent_wrench_composer.composed_torque,
+                    body_ids=self._ALL_BODY_INDICES_WP,
+                    env_ids=self._ALL_INDICES_WP,
+                )
+                # Apply both instantaneous and permanent wrench to the simulation
                 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),
+                    force_data=self._instantaneous_wrench_composer.composed_force_as_torch.view(-1, 3),
+                    torque_data=self._instantaneous_wrench_composer.composed_torque_as_torch.view(-1, 3),
+                    position_data=None,
                     indices=self._ALL_INDICES,
-                    is_global=self._use_global_wrench_frame,
+                    is_global=False,
                 )
             else:
+                # Apply permanent wrench to the simulation
                 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),
+                    force_data=self._permanent_wrench_composer.composed_force_as_torch.view(-1, 3),
+                    torque_data=self._permanent_wrench_composer.composed_torque_as_torch.view(-1, 3),
                     position_data=None,
                     indices=self._ALL_INDICES,
-                    is_global=self._use_global_wrench_frame,
+                    is_global=False,
                 )
+        self._instantaneous_wrench_composer.reset()
 
         # apply actuator models
         self._apply_actuator_model()
@@ -985,18 +1025,6 @@ def set_external_force_and_torque(
                 # example of disabling external wrench
                 asset.set_external_force_and_torque(forces=torch.zeros(0, 3), torques=torch.zeros(0, 3))
 
-        .. caution::
-            If the function is called consecutively with and with different values for ``is_global``, then the
-            all the external wrenches will be applied in the frame specified by the last call.
-
-            .. code-block:: python
-
-                # example of setting external wrench in the global frame
-                asset.set_external_force_and_torque(forces=torch.ones(1, 1, 3), env_ids=[0], is_global=True)
-                # example of setting external wrench in the link frame
-                asset.set_external_force_and_torque(forces=torch.ones(1, 1, 3), env_ids=[1], is_global=False)
-                # Both environments will have the external wrenches applied in the link frame
-
         .. note::
             This function does not apply the external wrench to the simulation. It only fills the buffers with
             the desired values. To apply the external wrench, call the :meth:`write_data_to_sim` function
@@ -1011,53 +1039,42 @@ def set_external_force_and_torque(
             is_global: Whether to apply the external wrench in the global frame. Defaults to False. If set to False,
                 the external wrench is applied in the link frame of the articulations' bodies.
         """
-        if forces.any() or torques.any():
-            self.has_external_wrench = True
-        else:
-            self.has_external_wrench = False
+        logger.warning(
+            "The function 'set_external_force_and_torque' will be deprecated in a future release. Please"
+            " use 'permanent_wrench_composer.set_forces_and_torques' instead."
+        )
+        if forces is None and torques is None:
+            logger.warning("No forces or torques provided. No permanent external wrench will be applied.")
 
         # resolve all indices
         # -- env_ids
         if env_ids is None:
-            env_ids = self._ALL_INDICES
+            env_ids = self._ALL_INDICES_WP
         elif not isinstance(env_ids, torch.Tensor):
-            env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device)
+            env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device)
+        else:
+            env_ids = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32)
         # -- body_ids
         if body_ids is None:
-            body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)
+            body_ids = self._ALL_BODY_INDICES_WP
         elif isinstance(body_ids, slice):
-            body_ids = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)[body_ids]
-        elif not isinstance(body_ids, torch.Tensor):
-            body_ids = torch.tensor(body_ids, dtype=torch.long, device=self.device)
-
-        # note: we need to do this complicated indexing since torch doesn't support multi-indexing
-        # create global body indices from env_ids and env_body_ids
-        # (env_id * total_bodies_per_env) + body_id
-        indices = body_ids.repeat(len(env_ids), 1) + env_ids.unsqueeze(1) * self.num_bodies
-        indices = indices.view(-1)
-        # set into internal buffers
-        # note: these are applied in the write_to_sim function
-        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 is_global != self._use_global_wrench_frame:
-            logger.warning(
-                f"The external wrench frame has been changed from {self._use_global_wrench_frame} to {is_global}. This"
-                " may lead to unexpected behavior."
+            body_ids = wp.from_torch(
+                torch.arange(self.num_bodies, dtype=torch.int32, device=self.device)[body_ids], dtype=wp.int32
             )
-            self._use_global_wrench_frame = is_global
-
-        # 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)
+        elif not isinstance(body_ids, torch.Tensor):
+            body_ids = wp.array(body_ids, dtype=wp.int32, device=self.device)
         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
+            body_ids = wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32)
+
+        # Write to wrench composer
+        self._permanent_wrench_composer.set_forces_and_torques(
+            forces=wp.from_torch(forces, dtype=wp.vec3f) if forces is not None else None,
+            torques=wp.from_torch(torques, dtype=wp.vec3f) if torques is not None else None,
+            positions=wp.from_torch(positions, dtype=wp.vec3f) if positions is not None else None,
+            body_ids=body_ids,
+            env_ids=env_ids,
+            is_global=is_global,
+        )
 
     def set_joint_position_target(
         self, target: torch.Tensor, joint_ids: Sequence[int] | slice | None = None, env_ids: Sequence[int] | None = None
@@ -1576,14 +1593,13 @@ def _initialize_impl(self):
     def _create_buffers(self):
         # constants
         self._ALL_INDICES = torch.arange(self.num_instances, dtype=torch.long, device=self.device)
+        self._ALL_BODY_INDICES = torch.arange(self.num_bodies, dtype=torch.long, device=self.device)
+        self._ALL_INDICES_WP = wp.from_torch(self._ALL_INDICES.to(torch.int32), dtype=wp.int32)
+        self._ALL_BODY_INDICES_WP = wp.from_torch(self._ALL_BODY_INDICES.to(torch.int32), dtype=wp.int32)
 
-        # 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)
-        self._use_global_wrench_frame = False
+        # external wrench composer
+        self._instantaneous_wrench_composer = WrenchComposer(self)
+        self._permanent_wrench_composer = WrenchComposer(self)
 
         # asset named data
         self._data.joint_names = self.joint_names