diff --git a/source/isaaclab/docs/CHANGELOG.rst b/source/isaaclab/docs/CHANGELOG.rst index fd1ecc9eeed..6a1cddfd0cf 100644 --- a/source/isaaclab/docs/CHANGELOG.rst +++ b/source/isaaclab/docs/CHANGELOG.rst @@ -4,12 +4,38 @@ Changelog 0.54.4 (2026-02-13) ~~~~~~~~~~~~~~~~~~~ +Changed +^^^^^^^ + +* Refactored :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to a dual-buffer architecture with separate + global (world-frame) and local (body-frame) buffers. A new + :meth:`~isaaclab.utils.wrench_composer.WrenchComposer.compose_to_body_frame` method rotates global forces/torques into + the body frame at apply time using the current body orientation, then sums with local forces/torques. Composed + wrenches are now applied to PhysX with ``is_global=False``. +* Replaced the ``add_forces_and_torques_at_position`` and ``set_forces_and_torques_at_position`` warp kernels with + ``set_forces_to_dual_buffers``, ``add_forces_to_dual_buffers``, ``add_raw_wrench_buffers``, and + ``compose_wrench_to_body_frame``. + Fixed ^^^^^ -* Fixed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` to correctly handle composed wrenches when link poses - change during simulation. Forces and torques are composed and stored in "mixed" frame representation (global frame - orientation, link frame position) and passed to Physx with `is_global=True`. +* Fixed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` not correctly updating the composed torque from global + positional forces when the body moves. Global positional torques are now stored as ``cross(P, F)`` about the world + origin and corrected at compose time via ``-cross(link_pos, F)`` to produce torque about the current CoM. +* Fixed :meth:`~isaaclab.utils.wrench_composer.WrenchComposer.reset` not clearing the ``_active`` flag when called + with ``slice(None)`` (the path taken by all asset reset methods). Added a ``ValueError`` for unsupported arbitrary + slices. +* Fixed :class:`~isaaclab.utils.wrench_composer.WrenchComposer` producing spurious torque when global forces are + applied without explicit positions. Previously, the compose kernel always applied the ``-cross(link_pos, F)`` + correction, causing torque proportional to the body's distance from the world origin. Global forces without + positions are now routed to a separate ``global_force_at_com_w`` buffer that bypasses the positional torque + correction, correctly applying the force at the body's center of mass. + +Added +^^^^^ + +* Added integration tests for :class:`~isaaclab.utils.wrench_composer.WrenchComposer` that validate global and local + force/torque behavior with a full physics simulation in the loop. 0.54.3 (2026-02-04) diff --git a/source/isaaclab/isaaclab/assets/articulation/articulation.py b/source/isaaclab/isaaclab/assets/articulation/articulation.py index de19097e3c1..d92e5c29c9e 100644 --- a/source/isaaclab/isaaclab/assets/articulation/articulation.py +++ b/source/isaaclab/isaaclab/assets/articulation/articulation.py @@ -228,30 +228,23 @@ def write_data_to_sim(self): # write external wrench 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, - is_global=True, - ) - # Apply both instantaneous and permanent wrench to the simulation + self._instantaneous_wrench_composer.add_raw_buffers_from(self._permanent_wrench_composer) + self._instantaneous_wrench_composer.compose_to_body_frame() self.root_physx_view.apply_forces_and_torques_at_position( - 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), + force_data=self._instantaneous_wrench_composer.out_force_b_as_torch.view(-1, 3), + torque_data=self._instantaneous_wrench_composer.out_torque_b_as_torch.view(-1, 3), position_data=None, indices=self._ALL_INDICES, - is_global=True, + is_global=False, ) else: - # Apply permanent wrench to the simulation + self._permanent_wrench_composer.compose_to_body_frame() self.root_physx_view.apply_forces_and_torques_at_position( - 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), + force_data=self._permanent_wrench_composer.out_force_b_as_torch.view(-1, 3), + torque_data=self._permanent_wrench_composer.out_torque_b_as_torch.view(-1, 3), position_data=None, indices=self._ALL_INDICES, - is_global=True, + is_global=False, ) self._instantaneous_wrench_composer.reset() @@ -1018,13 +1011,9 @@ def set_external_force_and_torque( 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 - of external wrench to the simulation. - - .. code-block:: python - - # example of disabling external wrench - asset.set_external_force_and_torque(forces=torch.zeros(0, 3), torques=torch.zeros(0, 3)) + This method clears **all** internal wrench buffers before writing the provided values. + Any previously set forces or torques (local or global) are discarded. To accumulate + on top of existing values, use ``permanent_wrench_composer.add_forces_and_torques`` instead. .. note:: This function does not apply the external wrench to the simulation. It only fills the buffers with @@ -1032,13 +1021,20 @@ def set_external_force_and_torque( right before the simulation step. 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. + forces: External forces. Shape is (len(env_ids), len(body_ids), 3). + When ``is_global=False``, forces are in the bodies' local frame. + When ``is_global=True``, forces are in the world frame. + torques: External torques. Shape is (len(env_ids), len(body_ids), 3). + When ``is_global=False``, torques are in the bodies' local frame. + When ``is_global=True``, torques are in the world frame. + positions: Application points for forces. Shape is (len(env_ids), len(body_ids), 3). + Defaults to None. + When ``is_global=False``, positions are local offsets from the link frame. + When ``is_global=True``, positions are world-frame coordinates. If None, + forces are applied at the body's center of mass (no positional torque). 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). - 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. + is_global: Whether forces and torques are in the global (world) frame. Defaults to False. """ logger.warning( "The function 'set_external_force_and_torque' will be deprecated in a future release. Please" diff --git a/source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py index 52789ff2a28..869459c554b 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py +++ b/source/isaaclab/isaaclab/assets/rigid_object/rigid_object.py @@ -141,30 +141,23 @@ def write_data_to_sim(self): # write external wrench 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, - is_global=True, - ) - # Apply both instantaneous and permanent wrench to the simulation + self._instantaneous_wrench_composer.add_raw_buffers_from(self._permanent_wrench_composer) + self._instantaneous_wrench_composer.compose_to_body_frame() self.root_physx_view.apply_forces_and_torques_at_position( - 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), + force_data=self._instantaneous_wrench_composer.out_force_b_as_torch.view(-1, 3), + torque_data=self._instantaneous_wrench_composer.out_torque_b_as_torch.view(-1, 3), position_data=None, indices=self._ALL_INDICES, - is_global=True, + is_global=False, ) else: - # Apply permanent wrench to the simulation + self._permanent_wrench_composer.compose_to_body_frame() self.root_physx_view.apply_forces_and_torques_at_position( - 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), + force_data=self._permanent_wrench_composer.out_force_b_as_torch.view(-1, 3), + torque_data=self._permanent_wrench_composer.out_torque_b_as_torch.view(-1, 3), position_data=None, indices=self._ALL_INDICES, - is_global=True, + is_global=False, ) self._instantaneous_wrench_composer.reset() @@ -417,13 +410,9 @@ def set_external_force_and_torque( 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 - of external wrench to the simulation. - - .. code-block:: python - - # example of disabling external wrench - asset.set_external_force_and_torque(forces=torch.zeros(0, 3), torques=torch.zeros(0, 3)) + This method clears **all** internal wrench buffers before writing the provided values. + Any previously set forces or torques (local or global) are discarded. To accumulate + on top of existing values, use ``permanent_wrench_composer.add_forces_and_torques`` instead. .. note:: This function does not apply the external wrench to the simulation. It only fills the buffers with @@ -431,14 +420,20 @@ def set_external_force_and_torque( right before the simulation step. 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). + forces: External forces. Shape is (len(env_ids), len(body_ids), 3). + When ``is_global=False``, forces are in the bodies' local frame. + When ``is_global=True``, forces are in the world frame. + torques: External torques. Shape is (len(env_ids), len(body_ids), 3). + When ``is_global=False``, torques are in the bodies' local frame. + When ``is_global=True``, torques are in the world frame. + positions: Application points for forces. Shape is (len(env_ids), len(body_ids), 3). Defaults to None. + When ``is_global=False``, positions are local offsets from the link frame. + When ``is_global=True``, positions are world-frame coordinates. If None, + forces are applied at the body's center of mass (no positional torque). 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). - 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 bodies. + is_global: Whether forces and torques are in the global (world) frame. Defaults to False. """ logger.warning( "The function 'set_external_force_and_torque' will be deprecated in a future release. Please" diff --git a/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection.py b/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection.py index c2e0f2fbfc3..98d98382ebc 100644 --- a/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection.py +++ b/source/isaaclab/isaaclab/assets/rigid_object_collection/rigid_object_collection.py @@ -185,30 +185,23 @@ def write_data_to_sim(self): # write external wrench 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_OBJ_INDICES_WP, - env_ids=self._ALL_ENV_INDICES_WP, - is_global=True, - ) - # Apply both instantaneous and permanent wrench to the simulation + self._instantaneous_wrench_composer.add_raw_buffers_from(self._permanent_wrench_composer) + self._instantaneous_wrench_composer.compose_to_body_frame() self.root_physx_view.apply_forces_and_torques_at_position( - force_data=self.reshape_data_to_view(self._instantaneous_wrench_composer.composed_force_as_torch), - torque_data=self.reshape_data_to_view(self._instantaneous_wrench_composer.composed_torque_as_torch), + force_data=self.reshape_data_to_view(self._instantaneous_wrench_composer.out_force_b_as_torch), + torque_data=self.reshape_data_to_view(self._instantaneous_wrench_composer.out_torque_b_as_torch), position_data=None, indices=self._env_obj_ids_to_view_ids(self._ALL_ENV_INDICES, self._ALL_OBJ_INDICES), - is_global=True, + is_global=False, ) else: - # Apply permanent wrench to the simulation + self._permanent_wrench_composer.compose_to_body_frame() self.root_physx_view.apply_forces_and_torques_at_position( - force_data=self.reshape_data_to_view(self._permanent_wrench_composer.composed_force_as_torch), - torque_data=self.reshape_data_to_view(self._permanent_wrench_composer.composed_torque_as_torch), + force_data=self.reshape_data_to_view(self._permanent_wrench_composer.out_force_b_as_torch), + torque_data=self.reshape_data_to_view(self._permanent_wrench_composer.out_torque_b_as_torch), position_data=None, indices=self._env_obj_ids_to_view_ids(self._ALL_ENV_INDICES, self._ALL_OBJ_INDICES), - is_global=True, + is_global=False, ) self._instantaneous_wrench_composer.reset() @@ -527,13 +520,9 @@ def set_external_force_and_torque( into buffers which are then applied to the simulation at every step. .. caution:: - If the function is called with empty forces and torques, then this function disables the application - of external wrench to the simulation. - - .. code-block:: python - - # example of disabling external wrench - asset.set_external_force_and_torque(forces=torch.zeros(0, 0, 3), torques=torch.zeros(0, 0, 3)) + This method clears **all** internal wrench buffers before writing the provided values. + Any previously set forces or torques (local or global) are discarded. To accumulate + on top of existing values, use ``permanent_wrench_composer.add_forces_and_torques`` instead. .. note:: This function does not apply the external wrench to the simulation. It only fills the buffers with @@ -541,13 +530,20 @@ def set_external_force_and_torque( right before the simulation step. 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). + forces: External forces. Shape is (len(env_ids), len(object_ids), 3). + When ``is_global=False``, forces are in the bodies' local frame. + When ``is_global=True``, forces are in the world frame. + torques: External torques. Shape is (len(env_ids), len(object_ids), 3). + When ``is_global=False``, torques are in the bodies' local frame. + When ``is_global=True``, torques are in the world frame. + positions: Application points for forces. Shape is (len(env_ids), len(object_ids), 3). + Defaults to None. + When ``is_global=False``, positions are local offsets from the link frame. + When ``is_global=True``, positions are world-frame coordinates. If None, + forces are applied at the body's center of mass (no positional torque). 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). - 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 bodies. + is_global: Whether forces and torques are in the global (world) frame. Defaults to False. """ logger.warning( "The function 'set_external_force_and_torque' will be deprecated in a future release. Please" diff --git a/source/isaaclab/isaaclab/utils/warp/kernels.py b/source/isaaclab/isaaclab/utils/warp/kernels.py index 0ad821fcf64..d26f4c1df0f 100644 --- a/source/isaaclab/isaaclab/utils/warp/kernels.py +++ b/source/isaaclab/isaaclab/utils/warp/kernels.py @@ -306,180 +306,212 @@ def reshape_tiled_image( ## -@wp.func -def cast_position_to_mixed_frame( - position: wp.vec3f, link_position: wp.vec3f, link_quat: wp.quatf, is_global: bool -) -> wp.vec3f: - """Casts a position to the mixed frame. +@wp.kernel +def set_forces_to_dual_buffers( + env_ids: wp.array(dtype=wp.int32), + body_ids: wp.array(dtype=wp.int32), + forces: wp.array2d(dtype=wp.vec3f), + torques: wp.array2d(dtype=wp.vec3f), + positions: wp.array2d(dtype=wp.vec3f), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), + is_global: bool, +): + """Sets forces and torques into the appropriate global or local buffer. - Args: - position: The position to cast. - link_position: The link frame position. - link_quat: The link frame quaternion. - is_global: Whether the position is in the global frame. + Routes forces/torques to either global (world-frame) or local (body-frame) buffers + based on the ``is_global`` flag. Position offsets contribute torque via cross product. - Returns: - The position in the mixed frame. - """ - if is_global: - return position - link_position - else: - return wp.quat_rotate(link_quat, position) + For global forces with positions, the torque is stored as ``cross(P, F)`` (torque about + world origin). At compose time, the correction ``-cross(link_pos, F)`` converts this + to torque about the current CoM. - -@wp.func -def cast_force_to_mixed_frame(force: wp.vec3f, link_quat: wp.quatf, is_global: bool) -> wp.vec3f: - """Casts a force to the mixed frame. + Global forces without positions are routed to ``global_force_at_com_w`` — they are applied + at the body's CoM and produce no positional torque. Args: - force: The force to cast. - link_quat: The link frame quaternion. - is_global: Whether the force is applied in the global frame. - Returns: - The force in the link frame of the body. + env_ids: Environment indices. + body_ids: Body indices. + forces: Input forces. Shape: (len(env_ids), len(body_ids)). + torques: Input torques. Shape: (len(env_ids), len(body_ids)). + positions: Application positions. Shape: (len(env_ids), len(body_ids)). + global_force_w: Output global force buffer (positional). Shape: (num_envs, num_bodies). + global_torque_w: Output global torque buffer. Shape: (num_envs, num_bodies). + global_force_at_com_w: Output global force at CoM buffer (no torque). Shape: (num_envs, num_bodies). + local_force_b: Output local force buffer. Shape: (num_envs, num_bodies). + local_torque_b: Output local torque buffer. Shape: (num_envs, num_bodies). + is_global: Whether forces/torques are in the global frame. """ - if is_global: - return force - else: - return wp.quat_rotate(link_quat, force) - - -@wp.func -def cast_torque_to_mixed_frame(torque: wp.vec3f, link_quat: wp.quatf, is_global: bool) -> wp.vec3f: - """Casts a torque to the mixed frame. - - Args: - torque: The torque to cast. - link_quat: The link frame quaternion. - is_global: Whether the torque is applied in the global frame. + tid_env, tid_body = wp.tid() + ei = env_ids[tid_env] + bi = body_ids[tid_body] - Returns: - The torque in the mixed frame. - """ if is_global: - return torque + if torques: + global_torque_w[ei, bi] = torques[tid_env, tid_body] + if forces: + if positions: + global_force_w[ei, bi] = forces[tid_env, tid_body] + if torques: + global_torque_w[ei, bi] = global_torque_w[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + global_torque_w[ei, bi] = wp.cross(positions[tid_env, tid_body], forces[tid_env, tid_body]) + else: + global_force_at_com_w[ei, bi] = forces[tid_env, tid_body] else: - return wp.quat_rotate(link_quat, torque) + if torques: + local_torque_b[ei, bi] = torques[tid_env, tid_body] + if forces: + local_force_b[ei, bi] = forces[tid_env, tid_body] + if positions: + if torques: + local_torque_b[ei, bi] = local_torque_b[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + local_torque_b[ei, bi] = wp.cross(positions[tid_env, tid_body], forces[tid_env, tid_body]) @wp.kernel -def add_forces_and_torques_at_position( +def add_forces_to_dual_buffers( env_ids: wp.array(dtype=wp.int32), body_ids: wp.array(dtype=wp.int32), forces: wp.array2d(dtype=wp.vec3f), torques: wp.array2d(dtype=wp.vec3f), positions: wp.array2d(dtype=wp.vec3f), - link_positions: wp.array2d(dtype=wp.vec3f), - link_quaternions: wp.array2d(dtype=wp.quatf), - composed_forces_m: wp.array2d(dtype=wp.vec3f), - composed_torques_m: wp.array2d(dtype=wp.vec3f), + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), is_global: bool, ): - """Adds forces and torques to the composed force and torque at the user-provided positions. - When is_global is False, the user-provided positions are offsetting the application of the force relatively to the - link frame of the body. When is_global is True, the user-provided positions are the global positions of the force - application. + """Adds forces and torques into the appropriate global or local buffer. + + Same as :func:`set_forces_to_dual_buffers` but uses ``+=`` instead of ``=``. Args: - env_ids: The environment ids. - body_ids: The body ids. - forces: The forces. - torques: The torques. - positions: The positions. - link_positions: The link frame positions. - link_quaternions: The link frame quaternions. - composed_forces_m: The composed forces in mixed representation (origin at link frame, orientation in global frame). - composed_torques_m: The composed torques in mixed representation (origin at link frame, orientation in global frame). - is_global: Whether the forces and torques are in the global frame. + env_ids: Environment indices. + body_ids: Body indices. + forces: Input forces. Shape: (len(env_ids), len(body_ids)). + torques: Input torques. Shape: (len(env_ids), len(body_ids)). + positions: Application positions. Shape: (len(env_ids), len(body_ids)). + global_force_w: Output global force buffer (positional). Shape: (num_envs, num_bodies). + global_torque_w: Output global torque buffer. Shape: (num_envs, num_bodies). + global_force_at_com_w: Output global force at CoM buffer (no torque). Shape: (num_envs, num_bodies). + local_force_b: Output local force buffer. Shape: (num_envs, num_bodies). + local_torque_b: Output local torque buffer. Shape: (num_envs, num_bodies). + is_global: Whether forces/torques are in the global frame. """ - # get the thread id tid_env, tid_body = wp.tid() + ei = env_ids[tid_env] + bi = body_ids[tid_body] - # add the forces to the composed force, if the positions are provided, also adds a torque to the composed torque. - if forces: - # Convert forces to mixed frame and add to composed force - composed_forces_m[env_ids[tid_env], body_ids[tid_body]] += cast_force_to_mixed_frame( - forces[tid_env, tid_body], link_quaternions[env_ids[tid_env], body_ids[tid_body]], is_global - ) - # If position offset provided, add torque contribution: τ = (pos - link_origin) × force - if positions: - composed_torques_m[env_ids[tid_env], body_ids[tid_body]] += wp.cross( - cast_position_to_mixed_frame( - positions[tid_env, tid_body], - link_positions[env_ids[tid_env], body_ids[tid_body]], - link_quaternions[env_ids[tid_env], body_ids[tid_body]], - is_global, - ), - cast_force_to_mixed_frame( - forces[tid_env, tid_body], link_quaternions[env_ids[tid_env], body_ids[tid_body]], is_global - ), - ) - if torques: - composed_torques_m[env_ids[tid_env], body_ids[tid_body]] += cast_torque_to_mixed_frame( - torques[tid_env, tid_body], link_quaternions[env_ids[tid_env], body_ids[tid_body]], is_global - ) + if is_global: + if forces: + if positions: + global_force_w[ei, bi] = global_force_w[ei, bi] + forces[tid_env, tid_body] + global_torque_w[ei, bi] = global_torque_w[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + else: + global_force_at_com_w[ei, bi] = global_force_at_com_w[ei, bi] + forces[tid_env, tid_body] + if torques: + global_torque_w[ei, bi] = global_torque_w[ei, bi] + torques[tid_env, tid_body] + else: + if forces: + local_force_b[ei, bi] = local_force_b[ei, bi] + forces[tid_env, tid_body] + if positions: + local_torque_b[ei, bi] = local_torque_b[ei, bi] + wp.cross( + positions[tid_env, tid_body], forces[tid_env, tid_body] + ) + if torques: + local_torque_b[ei, bi] = local_torque_b[ei, bi] + torques[tid_env, tid_body] @wp.kernel -def set_forces_and_torques_at_position( - env_ids: wp.array(dtype=wp.int32), - body_ids: wp.array(dtype=wp.int32), - forces: wp.array2d(dtype=wp.vec3f), - torques: wp.array2d(dtype=wp.vec3f), - positions: wp.array2d(dtype=wp.vec3f), +def add_raw_wrench_buffers( + src_gf: wp.array2d(dtype=wp.vec3f), + src_gt: wp.array2d(dtype=wp.vec3f), + src_gfc: wp.array2d(dtype=wp.vec3f), + src_lf: wp.array2d(dtype=wp.vec3f), + src_lt: wp.array2d(dtype=wp.vec3f), + dst_gf: wp.array2d(dtype=wp.vec3f), + dst_gt: wp.array2d(dtype=wp.vec3f), + dst_gfc: wp.array2d(dtype=wp.vec3f), + dst_lf: wp.array2d(dtype=wp.vec3f), + dst_lt: wp.array2d(dtype=wp.vec3f), +): + """Element-wise adds source wrench buffers into destination buffers. + + Used to merge one composer's 5 buffers (global force/torque, global force at CoM, + local force/torque) into another composer's buffers. + + Args: + src_gf: Source global forces (positional). + src_gt: Source global torques. + src_gfc: Source global forces at CoM. + src_lf: Source local forces. + src_lt: Source local torques. + dst_gf: Destination global forces (modified in-place). + dst_gt: Destination global torques (modified in-place). + dst_gfc: Destination global forces at CoM (modified in-place). + dst_lf: Destination local forces (modified in-place). + dst_lt: Destination local torques (modified in-place). + """ + tid_env, tid_body = wp.tid() + dst_gf[tid_env, tid_body] = dst_gf[tid_env, tid_body] + src_gf[tid_env, tid_body] + dst_gt[tid_env, tid_body] = dst_gt[tid_env, tid_body] + src_gt[tid_env, tid_body] + dst_gfc[tid_env, tid_body] = dst_gfc[tid_env, tid_body] + src_gfc[tid_env, tid_body] + dst_lf[tid_env, tid_body] = dst_lf[tid_env, tid_body] + src_lf[tid_env, tid_body] + dst_lt[tid_env, tid_body] = dst_lt[tid_env, tid_body] + src_lt[tid_env, tid_body] + + +@wp.kernel +def compose_wrench_to_body_frame( + global_force_w: wp.array2d(dtype=wp.vec3f), + global_torque_w: wp.array2d(dtype=wp.vec3f), + global_force_at_com_w: wp.array2d(dtype=wp.vec3f), + local_force_b: wp.array2d(dtype=wp.vec3f), + local_torque_b: wp.array2d(dtype=wp.vec3f), link_positions: wp.array2d(dtype=wp.vec3f), link_quaternions: wp.array2d(dtype=wp.quatf), - composed_forces_m: wp.array2d(dtype=wp.vec3f), - composed_torques_m: wp.array2d(dtype=wp.vec3f), - is_global: bool, + out_force_b: wp.array2d(dtype=wp.vec3f), + out_torque_b: wp.array2d(dtype=wp.vec3f), ): - """Sets forces and torques to the composed force and torque at the user-provided positions. - When is_global is False, the user-provided positions are offsetting the application of the force relatively - to the link frame of the body. When is_global is True, the user-provided positions are the global positions - of the force application. + """Composes global and local wrench buffers into a single body-frame output. - All forces and torques are stored in "mixed" representation: expressed in global (world) frame - orientation but referenced to the link origin. Positions are NOT stored - they are used to compute - the torque contribution from forces applied at offset positions. + Global torques are stored about the world origin (``cross(P, F)``). This kernel + corrects them to be about the current CoM by subtracting ``cross(link_pos, F)``, + then rotates into body frame using ``quat_rotate_inv`` and adds local-frame values. - When is_global is False, the user-provided positions are local offsets from the link frame origin. - When is_global is True, the user-provided positions are global coordinates. + Global forces at CoM (``global_force_at_com_w``) contribute to force output but + produce no positional torque — they bypass the ``cross(link_pos, F)`` correction. Args: - env_ids: The environment ids. - body_ids: The body ids. - forces: The forces. - torques: The torques. - positions: The positions. - link_positions: The link frame positions. - link_quaternions: The link frame quaternions. - composed_forces_m: The composed forces in mixed representation (origin at link frame, orientation in global frame). - composed_torques_m: The composed torques in mixed representation (origin at link frame, orientation in global frame). - is_global: Whether the forces and torques are in the global frame. + global_force_w: Global forces in world frame (with positions, participate in torque correction). + global_torque_w: Global torques in world frame (about world origin). + global_force_at_com_w: Global forces applied at CoM (no positional torque). + local_force_b: Local forces in body frame. + local_torque_b: Local torques in body frame. + link_positions: Current link positions in world frame. + link_quaternions: Body quaternions (xyzw convention for warp). + out_force_b: Output composed force in body frame. + out_torque_b: Output composed torque in body frame. """ - # get the thread id tid_env, tid_body = wp.tid() - - # Add user-provided torque - if torques: - composed_torques_m[env_ids[tid_env], body_ids[tid_body]] = cast_torque_to_mixed_frame( - torques[tid_env, tid_body], link_quaternions[env_ids[tid_env], body_ids[tid_body]], is_global - ) - # Set the forces to the composed force, if the positions are provided, adds a torque to the composed torque. - if forces: - composed_forces_m[env_ids[tid_env], body_ids[tid_body]] = cast_force_to_mixed_frame( - forces[tid_env, tid_body], link_quaternions[env_ids[tid_env], body_ids[tid_body]], is_global - ) - # if there is a position offset, set the torque from the force at that position. - if positions: - composed_torques_m[env_ids[tid_env], body_ids[tid_body]] += wp.cross( - cast_position_to_mixed_frame( - positions[tid_env, tid_body], - link_positions[env_ids[tid_env], body_ids[tid_body]], - link_quaternions[env_ids[tid_env], body_ids[tid_body]], - is_global, - ), - cast_force_to_mixed_frame( - forces[tid_env, tid_body], link_quaternions[env_ids[tid_env], body_ids[tid_body]], is_global - ), - ) + q = link_quaternions[tid_env, tid_body] + # Total world-frame force: positional + at-CoM + total_force_w = global_force_w[tid_env, tid_body] + global_force_at_com_w[tid_env, tid_body] + # Correct global torque: stored about world origin → about current CoM + # Only positional forces (global_force_w) participate in the correction + corrected_torque_w = global_torque_w[tid_env, tid_body] - wp.cross( + link_positions[tid_env, tid_body], global_force_w[tid_env, tid_body] + ) + out_force_b[tid_env, tid_body] = wp.quat_rotate_inv(q, total_force_w) + local_force_b[tid_env, tid_body] + out_torque_b[tid_env, tid_body] = wp.quat_rotate_inv(q, corrected_torque_w) + local_torque_b[tid_env, tid_body] diff --git a/source/isaaclab/isaaclab/utils/wrench_composer.py b/source/isaaclab/isaaclab/utils/wrench_composer.py index 46f64098cc4..6f82c1d1ddd 100644 --- a/source/isaaclab/isaaclab/utils/wrench_composer.py +++ b/source/isaaclab/isaaclab/utils/wrench_composer.py @@ -5,13 +5,19 @@ from __future__ import annotations +import warnings from typing import TYPE_CHECKING import torch import warp as wp from isaaclab.utils.math import convert_quat -from isaaclab.utils.warp.kernels import add_forces_and_torques_at_position, set_forces_and_torques_at_position +from isaaclab.utils.warp.kernels import ( + add_forces_to_dual_buffers, + add_raw_wrench_buffers, + compose_wrench_to_body_frame, + set_forces_to_dual_buffers, +) if TYPE_CHECKING: from isaaclab.assets import Articulation, RigidObject, RigidObjectCollection @@ -19,14 +25,18 @@ class WrenchComposer: def __init__(self, asset: Articulation | RigidObject | RigidObjectCollection) -> None: - """Wrench composer. + """Wrench composer with dual-buffer architecture (global + local). - This class composes forces and torques from multiple sources into a single wrench per body. - Forces and torques are stored in "mixed" representation: expressed in frame whose orientation is global, - while the origin is at the link origin. This allows for straightforward composition of forces and torques. + Forces and torques are stored in two separate pairs of buffers: + - **Global buffers** (``_global_force_w``, ``_global_torque_w``): world-frame forces/torques. + - **Local buffers** (``_local_force_b``, ``_local_torque_b``): body-frame forces/torques. + + At apply time, :meth:`compose_to_body_frame` rotates global forces into the body frame + using the current body quaternion and sums with local forces, producing a single + body-frame wrench that can be applied with ``is_global=False``. Args: - asset: Asset to use. Defaults to None. + asset: Asset to use. """ self.num_envs = asset.num_instances # Avoid isinstance to prevent circular import issues, use attribute presence instead. @@ -37,6 +47,7 @@ def __init__(self, asset: Articulation | RigidObject | RigidObjectCollection) -> self.device = asset.device self._asset = asset self._active = False + self._dirty = False # Avoid isinstance here due to potential circular import issues; check by attribute presence instead. if hasattr(self._asset.data, "body_link_pos_w") and hasattr(self._asset.data, "body_link_quat_w"): @@ -48,9 +59,20 @@ def __init__(self, asset: Articulation | RigidObject | RigidObjectCollection) -> else: raise ValueError(f"Unsupported asset type: {self._asset.__class__.__name__}") - # Create buffers - all forces and torques are stored in mixed representation: origin at link frame, orientation in global frame. - self._composed_force_m = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) - self._composed_torque_m = wp.zeros((self.num_envs, self.num_bodies), dtype=wp.vec3f, device=self.device) + shape = (self.num_envs, self.num_bodies) + + # Input buffers: global (world-frame) and local (body-frame) + self._global_force_w = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._global_torque_w = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._global_force_at_com_w = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._local_force_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._local_torque_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + + # Output buffers: composed body-frame wrench + self._out_force_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + self._out_torque_b = wp.zeros(shape, dtype=wp.vec3f, device=self.device) + + # Index arrays self._ALL_ENV_INDICES_WP = wp.from_torch( torch.arange(self.num_envs, dtype=torch.int32, device=self.device), dtype=wp.int32 ) @@ -58,12 +80,9 @@ def __init__(self, asset: Articulation | RigidObject | RigidObjectCollection) -> torch.arange(self.num_bodies, dtype=torch.int32, device=self.device), dtype=wp.int32 ) - # Pinning the composed force and torque to the torch tensor to avoid copying the data to the torch tensor - self._composed_force_m_torch = wp.to_torch(self._composed_force_m) - self._composed_torque_m_torch = wp.to_torch(self._composed_torque_m) - # Pinning the environment and body indices to the torch tensor to allow for slicing. - self._ALL_ENV_INDICES_TORCH = wp.to_torch(self._ALL_ENV_INDICES_WP) - self._ALL_BODY_INDICES_TORCH = wp.to_torch(self._ALL_BODY_INDICES_WP) + # Pinned torch views of output buffers (for PhysX apply calls) + self._out_force_b_torch = wp.to_torch(self._out_force_b) + self._out_torque_b_torch = wp.to_torch(self._out_torque_b) @property def active(self) -> bool: @@ -71,51 +90,69 @@ def active(self) -> bool: return self._active @property - def composed_force(self) -> wp.array: - """Composed force mixed representation: origin at link frame, orientation in global frame. + def global_force_w(self) -> wp.array: + """Global (world-frame) force buffer. Shape: (num_envs, num_bodies) vec3f.""" + return self._global_force_w - Forces are stored in "mixed" representation: global frame orientation, applied at link origin. - Any position offsets provided when setting forces contribute to torque, not to this force buffer. + @property + def global_torque_w(self) -> wp.array: + """Global (world-frame) torque buffer. Shape: (num_envs, num_bodies) vec3f.""" + return self._global_torque_w - Returns: - wp.array: Composed force in mixed representation. Shape: (num_envs, num_bodies, 3) - """ - return self._composed_force_m + @property + def global_force_at_com_w(self) -> wp.array: + """Global force applied at CoM (no positional torque). Shape: (num_envs, num_bodies) vec3f.""" + return self._global_force_at_com_w + + @property + def local_force_b(self) -> wp.array: + """Local (body-frame) force buffer. Shape: (num_envs, num_bodies) vec3f.""" + return self._local_force_b @property - def composed_torque(self) -> wp.array: - """Composed torque in mixed representation: origin at link frame, orientation in global frame. + def local_torque_b(self) -> wp.array: + """Local (body-frame) torque buffer. Shape: (num_envs, num_bodies) vec3f.""" + return self._local_torque_b - Torques are stored in "mixed" representation. - This includes both user-provided torques and torque contributions from forces applied at - offset positions (τ = (pos - link_origin) × force). + @property + def out_force_b(self) -> wp.array: + """Composed output force in body frame. Shape: (num_envs, num_bodies) vec3f. - Returns: - wp.array: Composed torque in global frame. Shape: (num_envs, num_bodies, 3) + If the output is stale (buffers were modified since last :meth:`compose_to_body_frame`), + this will automatically recompose and emit a performance warning. """ - return self._composed_torque_m + self._ensure_composed() + return self._out_force_b @property - def composed_force_as_torch(self) -> torch.Tensor: - """Composed force in mixed representation: origin at link frame, orientation in global frame, as torch tensor. + def out_torque_b(self) -> wp.array: + """Composed output torque in body frame. Shape: (num_envs, num_bodies) vec3f. - Forces are stored in "mixed" representation. - - Returns: - torch.Tensor: Composed force in mixed representation. Shape: (num_envs, num_bodies, 3) + If the output is stale (buffers were modified since last :meth:`compose_to_body_frame`), + this will automatically recompose and emit a performance warning. """ - return self._composed_force_m_torch + self._ensure_composed() + return self._out_torque_b @property - def composed_torque_as_torch(self) -> torch.Tensor: - """Composed torque mixed representation: origin at link frame, orientation in global frame, as torch tensor. + def out_force_b_as_torch(self) -> torch.Tensor: + """Composed output force in body frame as torch tensor. Shape: (num_envs, num_bodies, 3). + + If the output is stale (buffers were modified since last :meth:`compose_to_body_frame`), + this will automatically recompose and emit a performance warning. + """ + self._ensure_composed() + return self._out_force_b_torch - Torques are stored in "mixed" representation. + @property + def out_torque_b_as_torch(self) -> torch.Tensor: + """Composed output torque in body frame as torch tensor. Shape: (num_envs, num_bodies, 3). - Returns: - torch.Tensor: Composed torque in mixed representation. Shape: (num_envs, num_bodies, 3) + If the output is stale (buffers were modified since last :meth:`compose_to_body_frame`), + this will automatically recompose and emit a performance warning. """ - return self._composed_torque_m_torch + self._ensure_composed() + return self._out_torque_b_torch def add_forces_and_torques( self, @@ -126,63 +163,37 @@ def add_forces_and_torques( env_ids: wp.array | torch.Tensor | None = None, is_global: bool = False, ): - """Add forces and torques to the composed force and torque. + """Add forces and torques to the appropriate global or local buffer. - It can compose global wrenches and local wrenches. - It first convert them to the mixed representation and then add them to the already composed force and torque. + Routes to global buffers when ``is_global=True``, local buffers when ``is_global=False``. + Position offsets contribute additional torque via cross product. - Forces and torques are always stored in mixed representation: global frame orientation and application point at the link frame. + When ``is_global=True``: - Positions are NOT stored - they are used to compute torque contributions from forces applied at - offset positions (τ = (pos - link_origin) × force). + - Forces **with** positions are stored in the global positional buffer. The torque + ``cross(P, F)`` is stored about the world origin and corrected at compose time. + - Forces **without** positions are applied at the body's center of mass (no positional torque). + - Torques are stored directly in the global torque buffer. - .. note:: Users may want to call `reset` function after every simulation step to ensure no force is carried - over to the next step. However, this may not necessary if the user calls `set_forces_and_torques` function - instead of `add_forces_and_torques`. + When ``is_global=False``: + + - Forces and torques are stored in the local (body-frame) buffers. + - Positions are local offsets from the link frame contributing ``cross(pos, F)`` torque. Args: forces: Forces. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. torques: Torques. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. - positions: Positions. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. - When is_global is False, the user-provided positions are offsetting the application of the force relatively - to the link frame of the body. When is_global is True, the user-provided positions are the global positions - of the force application. + positions: Application points for forces. Shape: (len(env_ids), len(body_ids), 3). + Defaults to None. + When ``is_global=False``, positions are local offsets from the link frame. + When ``is_global=True``, positions are world-frame coordinates. If None, + forces are applied at the body's center of mass (no positional torque). body_ids: Body ids. Defaults to None (all bodies). env_ids: Environment ids. Defaults to None (all environments). is_global: Whether forces and torques are in global frame. Defaults to False. - - Raises: - ValueError: If the type of the input is not supported. - ValueError: If the input is a slice and it is not None. """ - # Resolve all indices - # -- env_ids - if env_ids is None: - env_ids = self._ALL_ENV_INDICES_WP - elif isinstance(env_ids, torch.Tensor): - env_ids = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) - elif isinstance(env_ids, list): - env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device) - elif isinstance(env_ids, slice): - if env_ids == slice(None): - env_ids = self._ALL_ENV_INDICES_WP - else: - raise ValueError(f"Doesn't support slice input for env_ids: {env_ids}") - # -- body_ids - if body_ids is None: - body_ids = self._ALL_BODY_INDICES_WP - elif isinstance(body_ids, torch.Tensor): - body_ids = wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32) - elif isinstance(body_ids, list): - body_ids = wp.array(body_ids, dtype=wp.int32, device=self.device) - elif isinstance(body_ids, slice): - if body_ids == slice(None): - body_ids = self._ALL_BODY_INDICES_WP - else: - raise ValueError(f"Doesn't support slice input for body_ids: {body_ids}") + env_ids, body_ids = self._resolve_indices(env_ids, body_ids) - # Resolve remaining inputs - # -- don't launch if no forces or torques are provided if forces is None and torques is None: return if isinstance(forces, torch.Tensor): @@ -192,17 +203,11 @@ def add_forces_and_torques( if isinstance(positions, torch.Tensor): positions = wp.from_torch(positions, dtype=wp.vec3f) - # Get the link positions and quaternions - self._link_positions = wp.from_torch(self._get_link_position_fn().clone(), dtype=wp.vec3f) - self._link_quaternions = wp.from_torch( - convert_quat(self._get_link_quaternion_fn().clone(), to="xyzw"), dtype=wp.quatf - ) - - # Set the active flag to true self._active = True + self._dirty = True wp.launch( - add_forces_and_torques_at_position, + add_forces_to_dual_buffers, dim=(env_ids.shape[0], body_ids.shape[0]), inputs=[ env_ids, @@ -210,10 +215,11 @@ def add_forces_and_torques( forces, torques, positions, - self._link_positions, - self._link_quaternions, - self._composed_force_m, - self._composed_torque_m, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, is_global, ], device=self.device, @@ -228,61 +234,40 @@ def set_forces_and_torques( env_ids: wp.array | torch.Tensor | None = None, is_global: bool = False, ): - """Set forces and torques to the composed force and torque (replaces existing values). + """Set forces and torques, replacing all existing values in every buffer. + + All 5 input buffers are cleared first, then the provided values are written to the + appropriate buffer. Use :meth:`add_forces_and_torques` to accumulate on top of existing values. - It can compose global wrenches and local wrenches. - It first convert them to the mixed representation and then add them to the already composed force and torque. + Routes to global buffers when ``is_global=True``, local buffers when ``is_global=False``. + Position offsets contribute additional torque via cross product. - Forces and torques are always stored in "mixed" representation: global frame orientation, - with application point at the link frame. + When ``is_global=True``: - Positions are NOT stored - they are used to compute torque - contributions from forces applied at offset positions (τ = (pos - link_origin) × force). + - Forces **with** positions are stored in the global positional buffer. The torque + ``cross(P, F)`` is stored about the world origin and corrected at compose time. + - Forces **without** positions are applied at the body's center of mass (no positional torque). + - Torques are stored directly in the global torque buffer. - The total torque set is: user_torque + cross(position - link_origin, force). + When ``is_global=False``: + + - Forces and torques are stored in the local (body-frame) buffers. + - Positions are local offsets from the link frame contributing ``cross(pos, F)`` torque. Args: forces: Forces. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. torques: Torques. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. - positions: Positions. Shape: (len(env_ids), len(body_ids), 3). Defaults to None. - When is_global is False, the user-provided positions are offsetting the application of the force relatively - to the link frame of the body. When is_global is True, the user-provided positions are the global positions - of the force application. + positions: Application points for forces. Shape: (len(env_ids), len(body_ids), 3). + Defaults to None. + When ``is_global=False``, positions are local offsets from the link frame. + When ``is_global=True``, positions are world-frame coordinates. If None, + forces are applied at the body's center of mass (no positional torque). body_ids: Body ids. Defaults to None (all bodies). env_ids: Environment ids. Defaults to None (all environments). is_global: Whether forces and torques are in global frame. Defaults to False. - - Raises: - ValueError: If the type of the input is not supported. - ValueError: If the input is a slice and it is not None. """ - # Resolve all indices - # -- env_ids - if env_ids is None: - env_ids = self._ALL_ENV_INDICES_WP - elif isinstance(env_ids, torch.Tensor): - env_ids = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) - elif isinstance(env_ids, list): - env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device) - elif isinstance(env_ids, slice): - if env_ids == slice(None): - env_ids = self._ALL_ENV_INDICES_WP - else: - raise ValueError(f"Doesn't support slice input for env_ids: {env_ids}") - # -- body_ids - if body_ids is None: - body_ids = self._ALL_BODY_INDICES_WP - elif isinstance(body_ids, torch.Tensor): - body_ids = wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32) - elif isinstance(body_ids, list): - body_ids = wp.array(body_ids, dtype=wp.int32, device=self.device) - elif isinstance(body_ids, slice): - if body_ids == slice(None): - body_ids = self._ALL_BODY_INDICES_WP - else: - raise ValueError(f"Doesn't support slice input for body_ids: {body_ids}") - # Resolve remaining inputs - # -- don't launch if no forces or torques are provided + env_ids, body_ids = self._resolve_indices(env_ids, body_ids) + if forces is None and torques is None: return if forces is None: @@ -298,17 +283,18 @@ def set_forces_and_torques( elif isinstance(positions, torch.Tensor): positions = wp.from_torch(positions, dtype=wp.vec3f) - # Get the link positions and quaternions - self._link_positions = wp.from_torch(self._get_link_position_fn().clone(), dtype=wp.vec3f) - self._link_quaternions = wp.from_torch( - convert_quat(self._get_link_quaternion_fn().clone(), to="xyzw"), dtype=wp.quatf - ) - - # Set the active flag to true self._active = True + self._dirty = True + + # Clear all input buffers first — set means "replace everything" + self._global_force_w.zero_() + self._global_torque_w.zero_() + self._global_force_at_com_w.zero_() + self._local_force_b.zero_() + self._local_torque_b.zero_() wp.launch( - set_forces_and_torques_at_position, + set_forces_to_dual_buffers, dim=(env_ids.shape[0], body_ids.shape[0]), inputs=[ env_ids, @@ -316,40 +302,153 @@ def set_forces_and_torques( forces, torques, positions, - self._link_positions, - self._link_quaternions, - self._composed_force_m, - self._composed_torque_m, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, is_global, ], device=self.device, ) - def reset(self, env_ids: wp.array | torch.Tensor | None = None): - """Reset the composed force and torque. + def add_raw_buffers_from(self, other: WrenchComposer): + """Element-wise add another composer's 4 input buffers into this composer's buffers. - This function will reset the composed force and torque to zero. - It will also make sure the link positions and quaternions are updated in the next call of the - `add_forces_and_torques` or `set_forces_and_torques` functions. + Args: + other: Source wrench composer whose buffers will be added. + """ + self._dirty = True + wp.launch( + add_raw_wrench_buffers, + dim=(self.num_envs, self.num_bodies), + inputs=[ + other._global_force_w, + other._global_torque_w, + other._global_force_at_com_w, + other._local_force_b, + other._local_torque_b, + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, + ], + device=self.device, + ) + + def compose_to_body_frame(self): + """Compose global and local buffers into body-frame output. - .. note:: This function should be called after every simulation step / reset to ensure no force is carried - over to the next step. + Fetches current link positions and quaternions. Global torques stored about + the world origin are corrected to be about the current CoM via + ``-cross(link_pos, global_force)``, then rotated into body frame via + ``quat_rotate_inv`` and summed with local forces/torques. Result is written + to :attr:`out_force_b` / :attr:`out_torque_b`. """ - if env_ids is None: - self._composed_force_m.zero_() - self._composed_torque_m.zero_() + link_quaternions = wp.from_torch( + convert_quat(self._get_link_quaternion_fn().clone(), to="xyzw"), dtype=wp.quatf + ) + link_positions = wp.from_torch(self._get_link_position_fn().clone(), dtype=wp.vec3f) + + wp.launch( + compose_wrench_to_body_frame, + dim=(self.num_envs, self.num_bodies), + inputs=[ + self._global_force_w, + self._global_torque_w, + self._global_force_at_com_w, + self._local_force_b, + self._local_torque_b, + link_positions, + link_quaternions, + self._out_force_b, + self._out_torque_b, + ], + device=self.device, + ) + self._dirty = False + + def reset(self, env_ids: wp.array | torch.Tensor | None = None, env_mask: wp.array | None = None): + """Reset all input and output buffers to zero. + + Args: + env_ids: Environment indices to reset. If None or slice(None), resets all. + env_mask: Environment mask (unused, kept for API compatibility). + + Raises: + ValueError: If env_ids is a slice other than slice(None). + """ + if env_ids is None or env_ids == slice(None): + self._global_force_w.zero_() + self._global_torque_w.zero_() + self._global_force_at_com_w.zero_() + self._local_force_b.zero_() + self._local_torque_b.zero_() + self._out_force_b.zero_() + self._out_torque_b.zero_() self._active = False else: + if isinstance(env_ids, slice): + raise ValueError( + f"WrenchComposer.reset() does not support arbitrary slices, got {env_ids!r}. " + "Use None, slice(None), or explicit index arrays instead." + ) indices = env_ids if isinstance(env_ids, torch.Tensor): indices = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) elif isinstance(env_ids, list): indices = wp.array(env_ids, dtype=wp.int32, device=self.device) - elif isinstance(env_ids, slice): - if env_ids == slice(None): - indices = self._ALL_ENV_INDICES_WP - else: - indices = env_ids - - self._composed_force_m[indices].zero_() - self._composed_torque_m[indices].zero_() + + self._global_force_w[indices].zero_() + self._global_torque_w[indices].zero_() + self._global_force_at_com_w[indices].zero_() + self._local_force_b[indices].zero_() + self._local_torque_b[indices].zero_() + self._out_force_b[indices].zero_() + self._out_torque_b[indices].zero_() + self._dirty = False + + def _ensure_composed(self): + """Ensure output buffers are up-to-date. If dirty, recomposes and warns.""" + if self._dirty: + warnings.warn( + "WrenchComposer: accessing output property triggered compose_to_body_frame() kernel launch. " + "Call compose_to_body_frame() explicitly before accessing output properties to avoid this overhead. " + "If you only need forces/torques in a single frame, use the raw buffer properties instead " + "(global_force_w, global_torque_w, local_force_b, local_torque_b) which require no composition.", + stacklevel=3, + ) + self.compose_to_body_frame() + + def _resolve_indices( + self, + env_ids: wp.array | torch.Tensor | None, + body_ids: wp.array | torch.Tensor | None, + ) -> tuple[wp.array, wp.array]: + """Resolve env and body indices to warp arrays.""" + if env_ids is None: + env_ids = self._ALL_ENV_INDICES_WP + elif isinstance(env_ids, torch.Tensor): + env_ids = wp.from_torch(env_ids.to(torch.int32), dtype=wp.int32) + elif isinstance(env_ids, list): + env_ids = wp.array(env_ids, dtype=wp.int32, device=self.device) + elif isinstance(env_ids, slice): + if env_ids == slice(None): + env_ids = self._ALL_ENV_INDICES_WP + else: + raise ValueError(f"Doesn't support slice input for env_ids: {env_ids}") + + if body_ids is None: + body_ids = self._ALL_BODY_INDICES_WP + elif isinstance(body_ids, torch.Tensor): + body_ids = wp.from_torch(body_ids.to(torch.int32), dtype=wp.int32) + elif isinstance(body_ids, list): + body_ids = wp.array(body_ids, dtype=wp.int32, device=self.device) + elif isinstance(body_ids, slice): + if body_ids == slice(None): + body_ids = self._ALL_BODY_INDICES_WP + else: + raise ValueError(f"Doesn't support slice input for body_ids: {body_ids}") + + return env_ids, body_ids diff --git a/source/isaaclab/test/assets/test_articulation.py b/source/isaaclab/test/assets/test_articulation.py index 9a983ab34c1..e55c13e947f 100644 --- a/source/isaaclab/test/assets/test_articulation.py +++ b/source/isaaclab/test/assets/test_articulation.py @@ -21,6 +21,7 @@ import pytest import torch +import warp as wp import isaaclab.sim as sim_utils import isaaclab.utils.math as math_utils @@ -836,8 +837,8 @@ def test_external_force_buffer(sim, num_articulations, device): # check if the articulation's force and torque buffers are correctly updated for i in range(num_articulations): - assert articulation.permanent_wrench_composer.composed_force_as_torch[i, 0, 0].item() == force - assert articulation.permanent_wrench_composer.composed_torque_as_torch[i, 0, 0].item() == force + assert wp.to_torch(articulation.permanent_wrench_composer.local_force_b)[i, 0, 0].item() == force + assert wp.to_torch(articulation.permanent_wrench_composer.local_torque_b)[i, 0, 0].item() == force # Check if the instantaneous wrench is correctly added to the permanent wrench articulation.instantaneous_wrench_composer.add_forces_and_torques( @@ -1551,10 +1552,10 @@ def test_reset(sim, num_articulations, device): # Reset should zero external forces and torques assert not articulation._instantaneous_wrench_composer.active assert not articulation._permanent_wrench_composer.active - assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force_as_torch) == 0 - assert torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque_as_torch) == 0 - assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force_as_torch) == 0 - assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque_as_torch) == 0 + assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.global_force_w)) == 0 + assert torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.local_force_b)) == 0 + assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.global_force_w)) == 0 + assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.local_force_b)) == 0 if num_articulations > 1: num_bodies = articulation.num_bodies @@ -1571,13 +1572,17 @@ def test_reset(sim, num_articulations, device): assert articulation._instantaneous_wrench_composer.active assert articulation._permanent_wrench_composer.active assert ( - torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_force_as_torch) == num_bodies * 3 + torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.local_force_b)) + == num_bodies * 3 ) assert ( - torch.count_nonzero(articulation._instantaneous_wrench_composer.composed_torque_as_torch) == num_bodies * 3 + torch.count_nonzero(wp.to_torch(articulation._instantaneous_wrench_composer.local_torque_b)) + == num_bodies * 3 + ) + assert torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.local_force_b)) == num_bodies * 3 + assert ( + torch.count_nonzero(wp.to_torch(articulation._permanent_wrench_composer.local_torque_b)) == num_bodies * 3 ) - assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_force_as_torch) == num_bodies * 3 - assert torch.count_nonzero(articulation._permanent_wrench_composer.composed_torque_as_torch) == num_bodies * 3 @pytest.mark.parametrize("num_articulations", [1, 2]) diff --git a/source/isaaclab/test/assets/test_rigid_object.py b/source/isaaclab/test/assets/test_rigid_object.py index 8de5361e29f..dbb7d7a7ed3 100644 --- a/source/isaaclab/test/assets/test_rigid_object.py +++ b/source/isaaclab/test/assets/test_rigid_object.py @@ -21,6 +21,7 @@ import pytest import torch +import warp as wp from flaky import flaky import isaaclab.sim as sim_utils @@ -249,8 +250,8 @@ def test_external_force_buffer(device): # check if the cube's force and torque buffers are correctly updated for i in range(cube_object.num_instances): - assert cube_object._permanent_wrench_composer.composed_force_as_torch[i, 0, 0].item() == force - assert cube_object._permanent_wrench_composer.composed_torque_as_torch[i, 0, 0].item() == force + assert wp.to_torch(cube_object._permanent_wrench_composer.local_force_b)[i, 0, 0].item() == force + assert wp.to_torch(cube_object._permanent_wrench_composer.local_torque_b)[i, 0, 0].item() == force # Check if the instantaneous wrench is correctly added to the permanent wrench cube_object.permanent_wrench_composer.add_forces_and_torques( @@ -419,14 +420,15 @@ def test_external_force_on_single_body_at_position(num_cubes, device): body_ids=body_ids, is_global=is_global, ) + cube_object._permanent_wrench_composer.compose_to_body_frame() torch.testing.assert_close( - cube_object._permanent_wrench_composer.composed_force_as_torch[:, 0, :], + cube_object._permanent_wrench_composer.out_force_b_as_torch[:, 0, :], desired_force[:, 0, :], rtol=1e-6, atol=1e-7, ) torch.testing.assert_close( - cube_object._permanent_wrench_composer.composed_torque_as_torch[:, 0, :], + cube_object._permanent_wrench_composer.out_torque_b_as_torch[:, 0, :], desired_torque[:, 0, :], rtol=1e-6, atol=1e-7, @@ -551,10 +553,10 @@ def test_reset_rigid_object(num_cubes, device): # Reset should zero external forces and torques assert not cube_object._instantaneous_wrench_composer.active assert not cube_object._permanent_wrench_composer.active - assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_force_as_torch) == 0 - assert torch.count_nonzero(cube_object._instantaneous_wrench_composer.composed_torque_as_torch) == 0 - assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_force_as_torch) == 0 - assert torch.count_nonzero(cube_object._permanent_wrench_composer.composed_torque_as_torch) == 0 + assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.global_force_w)) == 0 + assert torch.count_nonzero(wp.to_torch(cube_object._instantaneous_wrench_composer.local_force_b)) == 0 + assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.global_force_w)) == 0 + assert torch.count_nonzero(wp.to_torch(cube_object._permanent_wrench_composer.local_force_b)) == 0 @pytest.mark.parametrize("num_cubes", [1, 2]) diff --git a/source/isaaclab/test/assets/test_rigid_object_collection.py b/source/isaaclab/test/assets/test_rigid_object_collection.py index 8d919bf4d7a..bb17d782e67 100644 --- a/source/isaaclab/test/assets/test_rigid_object_collection.py +++ b/source/isaaclab/test/assets/test_rigid_object_collection.py @@ -20,6 +20,7 @@ import pytest import torch +import warp as wp import isaaclab.sim as sim_utils from isaaclab.assets import RigidObjectCfg, RigidObjectCollection, RigidObjectCollectionCfg @@ -252,8 +253,8 @@ def test_external_force_buffer(sim, device): # check if the object collection's force and torque buffers are correctly updated for i in range(num_envs): - assert object_collection._permanent_wrench_composer.composed_force_as_torch[i, 0, 0].item() == force - assert object_collection._permanent_wrench_composer.composed_torque_as_torch[i, 0, 0].item() == force + assert wp.to_torch(object_collection._permanent_wrench_composer.local_force_b)[i, 0, 0].item() == force + assert wp.to_torch(object_collection._permanent_wrench_composer.local_torque_b)[i, 0, 0].item() == force object_collection.instantaneous_wrench_composer.add_forces_and_torques( body_ids=object_ids, @@ -644,10 +645,12 @@ def test_reset_object_collection(sim, num_envs, num_cubes, device): # Reset should zero external forces and torques assert not object_collection._instantaneous_wrench_composer.active assert not object_collection._permanent_wrench_composer.active - assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_force_as_torch) == 0 - assert torch.count_nonzero(object_collection._instantaneous_wrench_composer.composed_torque_as_torch) == 0 - assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_force_as_torch) == 0 - assert torch.count_nonzero(object_collection._permanent_wrench_composer.composed_torque_as_torch) == 0 + assert ( + torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.global_force_w)) == 0 + ) + assert torch.count_nonzero(wp.to_torch(object_collection._instantaneous_wrench_composer.local_force_b)) == 0 + assert torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.global_force_w)) == 0 + assert torch.count_nonzero(wp.to_torch(object_collection._permanent_wrench_composer.local_force_b)) == 0 @pytest.mark.parametrize("num_envs", [1, 3]) diff --git a/source/isaaclab/test/utils/test_wrench_composer.py b/source/isaaclab/test/utils/test_wrench_composer.py index 8dfc56100a0..ce4b33f998e 100644 --- a/source/isaaclab/test/utils/test_wrench_composer.py +++ b/source/isaaclab/test/utils/test_wrench_composer.py @@ -8,6 +8,8 @@ # launch omniverse app simulation_app = AppLauncher(headless=True).app +import warnings as warnings_module + import numpy as np import pytest import torch @@ -150,105 +152,93 @@ def random_unit_quaternion_np(rng: np.random.Generator, shape: tuple) -> np.ndar RigidObject.register(MockRigidObject) +# ============================================================================ +# Basic Tests (identity quaternion, is_global=False by default) +# ============================================================================ + + @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100, 1000]) @pytest.mark.parametrize("num_bodies", [1, 3, 5, 10]) def test_wrench_composer_add_force(device: str, num_envs: int, num_bodies: int): - # Initialize random number generator + """Test adding local forces (default is_global=False) with identity quaternion.""" rng = np.random.default_rng(seed=0) for _ in range(10): mock_asset = MockRigidObject(num_envs, num_bodies, device) wrench_composer = WrenchComposer(mock_asset) - # Initialize hand-calculated composed force hand_calculated_composed_force_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) for _ in range(10): - # Get random number of envs and bodies and their indices num_envs_np = rng.integers(1, num_envs, endpoint=True) num_bodies_np = rng.integers(1, num_bodies, endpoint=True) env_ids_np = rng.choice(num_envs, size=num_envs_np, replace=False) body_ids_np = rng.choice(num_bodies, size=num_bodies_np, replace=False) - # Convert to warp arrays env_ids = wp.from_numpy(env_ids_np, dtype=wp.int32, device=device) body_ids = wp.from_numpy(body_ids_np, dtype=wp.int32, device=device) - # Get random forces forces_np = ( np.random.uniform(low=-100.0, high=100.0, size=(num_envs_np * num_bodies_np * 3)) .reshape(num_envs_np, num_bodies_np, 3) .astype(np.float32) ) forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) - # Add forces to wrench composer wrench_composer.add_forces_and_torques(forces=forces, body_ids=body_ids, env_ids=env_ids) - # Add forces to hand-calculated composed force hand_calculated_composed_force_np[env_ids_np[:, None], body_ids_np[None, :], :] += forces_np - # Get composed force from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) + # Local forces with identity quat → compose gives local forces unchanged + wrench_composer.compose_to_body_frame() + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100, 1000]) @pytest.mark.parametrize("num_bodies", [1, 3, 5, 10]) def test_wrench_composer_add_torque(device: str, num_envs: int, num_bodies: int): - # Initialize random number generator + """Test adding local torques (default is_global=False) with identity quaternion.""" rng = np.random.default_rng(seed=1) for _ in range(10): mock_asset = MockRigidObject(num_envs, num_bodies, device) wrench_composer = WrenchComposer(mock_asset) - # Initialize hand-calculated composed torque hand_calculated_composed_torque_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) for _ in range(10): - # Get random number of envs and bodies and their indices num_envs_np = rng.integers(1, num_envs, endpoint=True) num_bodies_np = rng.integers(1, num_bodies, endpoint=True) env_ids_np = rng.choice(num_envs, size=num_envs_np, replace=False) body_ids_np = rng.choice(num_bodies, size=num_bodies_np, replace=False) - # Convert to warp arrays env_ids = wp.from_numpy(env_ids_np, dtype=wp.int32, device=device) body_ids = wp.from_numpy(body_ids_np, dtype=wp.int32, device=device) - # Get random torques torques_np = ( np.random.uniform(low=-100.0, high=100.0, size=(num_envs_np * num_bodies_np * 3)) .reshape(num_envs_np, num_bodies_np, 3) .astype(np.float32) ) torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) - # Add torques to wrench composer wrench_composer.add_forces_and_torques(torques=torques, body_ids=body_ids, env_ids=env_ids) - # Add torques to hand-calculated composed torque hand_calculated_composed_torque_np[env_ids_np[:, None], body_ids_np[None, :], :] += torques_np - # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) + wrench_composer.compose_to_body_frame() + out_torque_np = wrench_composer.out_torque_b.numpy() + assert np.allclose(out_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100, 1000]) @pytest.mark.parametrize("num_bodies", [1, 3, 5, 10]) def test_add_forces_at_positons(device: str, num_envs: int, num_bodies: int): - """Test adding forces at local positions (offset from link frame).""" + """Test adding local forces at local positions (offset from link frame).""" rng = np.random.default_rng(seed=2) for _ in range(10): - # Initialize wrench composer mock_asset = MockRigidObject(num_envs, num_bodies, device) wrench_composer = WrenchComposer(mock_asset) - # Initialize hand-calculated composed force hand_calculated_composed_force_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) - # Initialize hand-calculated composed torque hand_calculated_composed_torque_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) for _ in range(10): - # Get random number of envs and bodies and their indices num_envs_np = rng.integers(1, num_envs, endpoint=True) num_bodies_np = rng.integers(1, num_bodies, endpoint=True) env_ids_np = rng.choice(num_envs, size=num_envs_np, replace=False) body_ids_np = rng.choice(num_bodies, size=num_bodies_np, replace=False) - # Convert to warp arrays env_ids = wp.from_numpy(env_ids_np, dtype=wp.int32, device=device) body_ids = wp.from_numpy(body_ids_np, dtype=wp.int32, device=device) - # Get random forces forces_np = ( np.random.uniform(low=-100.0, high=100.0, size=(num_envs_np * num_bodies_np * 3)) .reshape(num_envs_np, num_bodies_np, 3) @@ -261,47 +251,42 @@ def test_add_forces_at_positons(device: str, num_envs: int, num_bodies: int): ) forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) positions = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) - # Add forces at positions to wrench composer wrench_composer.add_forces_and_torques( forces=forces, positions=positions, body_ids=body_ids, env_ids=env_ids ) - # Add forces to hand-calculated composed force + # Local forces accumulate directly hand_calculated_composed_force_np[env_ids_np[:, None], body_ids_np[None, :], :] += forces_np - # Add torques to hand-calculated composed torque: torque = cross(position, force) + # Local torque from position: cross(local_pos, local_force) torques_from_forces = np.cross(positions_np, forces_np) for i in range(num_envs_np): for j in range(num_bodies_np): hand_calculated_composed_torque_np[env_ids_np[i], body_ids_np[j], :] += torques_from_forces[i, j, :] - # Get composed force from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) - # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) + wrench_composer.compose_to_body_frame() + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) + out_torque_np = wrench_composer.out_torque_b.numpy() + assert np.allclose(out_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100, 1000]) @pytest.mark.parametrize("num_bodies", [1, 3, 5, 10]) def test_add_torques_at_position(device: str, num_envs: int, num_bodies: int): + """Test that positions don't affect torque-only additions.""" rng = np.random.default_rng(seed=3) for _ in range(10): mock_asset = MockRigidObject(num_envs, num_bodies, device) wrench_composer = WrenchComposer(mock_asset) - # Initialize hand-calculated composed torque hand_calculated_composed_torque_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) for _ in range(10): - # Get random number of envs and bodies and their indices num_envs_np = rng.integers(1, num_envs, endpoint=True) num_bodies_np = rng.integers(1, num_bodies, endpoint=True) env_ids_np = rng.choice(num_envs, size=num_envs_np, replace=False) body_ids_np = rng.choice(num_bodies, size=num_bodies_np, replace=False) - # Convert to warp arrays env_ids = wp.from_numpy(env_ids_np, dtype=wp.int32, device=device) body_ids = wp.from_numpy(body_ids_np, dtype=wp.int32, device=device) - # Get random torques torques_np = ( np.random.uniform(low=-100.0, high=100.0, size=(num_envs_np * num_bodies_np * 3)) .reshape(num_envs_np, num_bodies_np, 3) @@ -314,40 +299,34 @@ def test_add_torques_at_position(device: str, num_envs: int, num_bodies: int): ) torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) positions = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) - # Add torques at positions to wrench composer wrench_composer.add_forces_and_torques( torques=torques, positions=positions, body_ids=body_ids, env_ids=env_ids ) - # Add torques to hand-calculated composed torque hand_calculated_composed_torque_np[env_ids_np[:, None], body_ids_np[None, :], :] += torques_np - # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) + wrench_composer.compose_to_body_frame() + out_torque_np = wrench_composer.out_torque_b.numpy() + assert np.allclose(out_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100, 1000]) @pytest.mark.parametrize("num_bodies", [1, 3, 5, 10]) def test_add_forces_and_torques_at_position(device: str, num_envs: int, num_bodies: int): - """Test adding forces and torques at local positions.""" + """Test adding local forces and torques at local positions.""" rng = np.random.default_rng(seed=4) for _ in range(10): mock_asset = MockRigidObject(num_envs, num_bodies, device) wrench_composer = WrenchComposer(mock_asset) - # Initialize hand-calculated composed force and torque hand_calculated_composed_force_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) hand_calculated_composed_torque_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) for _ in range(10): - # Get random number of envs and bodies and their indices num_envs_np = rng.integers(1, num_envs, endpoint=True) num_bodies_np = rng.integers(1, num_bodies, endpoint=True) env_ids_np = rng.choice(num_envs, size=num_envs_np, replace=False) body_ids_np = rng.choice(num_bodies, size=num_bodies_np, replace=False) - # Convert to warp arrays env_ids = wp.from_numpy(env_ids_np, dtype=wp.int32, device=device) body_ids = wp.from_numpy(body_ids_np, dtype=wp.int32, device=device) - # Get random forces and torques forces_np = ( np.random.uniform(low=-100.0, high=100.0, size=(num_envs_np * num_bodies_np * 3)) .reshape(num_envs_np, num_bodies_np, 3) @@ -366,43 +345,37 @@ def test_add_forces_and_torques_at_position(device: str, num_envs: int, num_bodi forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) positions = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) - # Add forces and torques at positions to wrench composer wrench_composer.add_forces_and_torques( forces=forces, torques=torques, positions=positions, body_ids=body_ids, env_ids=env_ids ) - # Add forces to hand-calculated composed force hand_calculated_composed_force_np[env_ids_np[:, None], body_ids_np[None, :], :] += forces_np - # Add torques to hand-calculated composed torque: torque = cross(position, force) + torque torques_from_forces = np.cross(positions_np, forces_np) for i in range(num_envs_np): for j in range(num_bodies_np): hand_calculated_composed_torque_np[env_ids_np[i], body_ids_np[j], :] += torques_from_forces[i, j, :] hand_calculated_composed_torque_np[env_ids_np[:, None], body_ids_np[None, :], :] += torques_np - # Get composed force from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) - # Get composed torque from wrench composer - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) + wrench_composer.compose_to_body_frame() + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, hand_calculated_composed_force_np, atol=1, rtol=1e-7) + out_torque_np = wrench_composer.out_torque_b.numpy() + assert np.allclose(out_torque_np, hand_calculated_composed_torque_np, atol=1, rtol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100, 1000]) @pytest.mark.parametrize("num_bodies", [1, 3, 5, 10]) def test_wrench_composer_reset(device: str, num_envs: int, num_bodies: int): + """Test that reset zeros all 4 input buffers and 2 output buffers.""" rng = np.random.default_rng(seed=5) for _ in range(10): mock_asset = MockRigidObject(num_envs, num_bodies, device) wrench_composer = WrenchComposer(mock_asset) - # Get random number of envs and bodies and their indices num_envs_np = rng.integers(1, num_envs, endpoint=True) num_bodies_np = rng.integers(1, num_bodies, endpoint=True) env_ids_np = rng.choice(num_envs, size=num_envs_np, replace=False) body_ids_np = rng.choice(num_bodies, size=num_bodies_np, replace=False) - # Convert to warp arrays env_ids = wp.from_numpy(env_ids_np, dtype=wp.int32, device=device) body_ids = wp.from_numpy(body_ids_np, dtype=wp.int32, device=device) - # Get random forces and torques forces_np = ( np.random.uniform(low=-100.0, high=100.0, size=(num_envs_np * num_bodies_np * 3)) .reshape(num_envs_np, num_bodies_np, 3) @@ -415,15 +388,26 @@ def test_wrench_composer_reset(device: str, num_envs: int, num_bodies: int): ) forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) - # Add forces and torques to wrench composer + # Add local forces/torques wrench_composer.add_forces_and_torques(forces=forces, torques=torques, body_ids=body_ids, env_ids=env_ids) - # Reset wrench composer + # Add global forces/torques + wrench_composer.add_forces_and_torques( + forces=forces, torques=torques, body_ids=body_ids, env_ids=env_ids, is_global=True + ) + # Compose to populate output buffers + wrench_composer.compose_to_body_frame() + # Reset wrench_composer.reset() - # Get composed force and torque from wrench composer - composed_force_np = wrench_composer.composed_force.numpy() - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_force_np, np.zeros((num_envs, num_bodies, 3)), atol=1, rtol=1e-7) - assert np.allclose(composed_torque_np, np.zeros((num_envs, num_bodies, 3)), atol=1, rtol=1e-7) + # All 7 buffers should be zero (5 input + 2 output) + zeros = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(wrench_composer.global_force_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_torque_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_force_at_com_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.local_force_b.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.local_torque_b.numpy(), zeros, atol=1e-7) + # Access _out_force_b directly to avoid triggering warning (dirty=False after reset) + assert np.allclose(wrench_composer._out_force_b.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer._out_torque_b.numpy(), zeros, atol=1e-7) # ============================================================================ @@ -434,164 +418,142 @@ def test_wrench_composer_reset(device: str, num_envs: int, num_bodies: int): @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100]) @pytest.mark.parametrize("num_bodies", [1, 3, 5]) -def test_global_forces_with_rotation(device: str, num_envs: int, num_bodies: int): - """Test that global forces stay unchanged in mixed representation (global orientation).""" +def test_global_forces_stored_in_global_buffer(device: str, num_envs: int, num_bodies: int): + """Test that global forces without positions are stored in the global_force_at_com_w buffer.""" rng = np.random.default_rng(seed=10) for _ in range(5): - # Create random link quaternions link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) link_quat_torch = torch.from_numpy(link_quat_np) - # Create mock asset with custom quaternions mock_asset = MockRigidObject(num_envs, num_bodies, device, link_quat=link_quat_torch) wrench_composer = WrenchComposer(mock_asset) - # Generate random global forces for all envs and bodies forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) - # Apply global forces wrench_composer.add_forces_and_torques(forces=forces_global, is_global=True) - # In mixed representation, global forces stay unchanged (already in global orientation) - expected_forces_mixed = forces_global_np - - # Verify - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, expected_forces_mixed, atol=1e-4, rtol=1e-5), ( - f"Global force in mixed repr failed.\nExpected:\n{expected_forces_mixed}\nGot:\n{composed_force_np}" - ) + # Global forces without positions stored in global_force_at_com_w buffer + global_force_at_com_np = wrench_composer.global_force_at_com_w.numpy() + assert np.allclose(global_force_at_com_np, forces_global_np, atol=1e-4, rtol=1e-5) + # Positional global_force_w should remain zero + global_force_np = wrench_composer.global_force_w.numpy() + assert np.allclose(global_force_np, np.zeros_like(forces_global_np), atol=1e-7) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 100]) @pytest.mark.parametrize("num_bodies", [1, 3, 5]) -def test_global_torques_with_rotation(device: str, num_envs: int, num_bodies: int): - """Test that global torques stay unchanged in mixed representation (global orientation).""" +def test_global_torques_stored_in_global_buffer(device: str, num_envs: int, num_bodies: int): + """Test that global torques are stored unchanged in the global buffer.""" rng = np.random.default_rng(seed=11) for _ in range(5): - # Create random link quaternions link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) link_quat_torch = torch.from_numpy(link_quat_np) - # Create mock asset with custom quaternions mock_asset = MockRigidObject(num_envs, num_bodies, device, link_quat=link_quat_torch) wrench_composer = WrenchComposer(mock_asset) - # Generate random global torques torques_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) torques_global = wp.from_numpy(torques_global_np, dtype=wp.vec3f, device=device) - # Apply global torques wrench_composer.add_forces_and_torques(torques=torques_global, is_global=True) - # In mixed representation, global torques stay unchanged (already in global orientation) - expected_torques_mixed = torques_global_np - - # Verify - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_torque_np, expected_torques_mixed, atol=1e-4, rtol=1e-5), ( - f"Global torque in mixed repr failed.\nExpected:\n{expected_torques_mixed}\nGot:\n{composed_torque_np}" - ) + global_torque_np = wrench_composer.global_torque_w.numpy() + assert np.allclose(global_torque_np, torques_global_np, atol=1e-4, rtol=1e-5) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) @pytest.mark.parametrize("num_envs", [1, 10, 50]) @pytest.mark.parametrize("num_bodies", [1, 3, 5]) def test_global_forces_at_global_position(device: str, num_envs: int, num_bodies: int): - """Test global forces at global positions with full coordinate transformation.""" + """Test global forces at global positions produce correct torque in global buffer. + + Global torque is stored about the world origin: cross(P, F). + After compose, the correction -cross(link_pos, F) gives torque about CoM: + cross(P, F) - cross(link_pos, F) = cross(P - link_pos, F). + """ rng = np.random.default_rng(seed=12) for _ in range(5): - # Create random link poses link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) link_pos_torch = torch.from_numpy(link_pos_np) link_quat_torch = torch.from_numpy(link_quat_np) - # Create mock asset mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) wrench_composer = WrenchComposer(mock_asset) - # Generate random global forces and positions forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) positions_global_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) positions_global = wp.from_numpy(positions_global_np, dtype=wp.vec3f, device=device) - # Apply global forces at global positions wrench_composer.add_forces_and_torques(forces=forces_global, positions=positions_global, is_global=True) - # Compute expected results in mixed representation: - # 1. Force stays unchanged (already in global orientation) - expected_forces_mixed = forces_global_np - - # 2. Position offset in global frame = global_position - link_position - position_offset_global = positions_global_np - link_pos_np + # Global force stored unchanged + expected_forces = forces_global_np + # Torque stored about world origin: cross(P, F) + expected_stored_torques = np.cross(positions_global_np, forces_global_np) - # 3. Torque = cross(position_offset, force) in global frame - expected_torques_mixed = np.cross(position_offset_global, forces_global_np) + global_force_np = wrench_composer.global_force_w.numpy() + assert np.allclose(global_force_np, expected_forces, atol=1e-3, rtol=1e-4) - # Verify forces - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, expected_forces_mixed, atol=1e-3, rtol=1e-4), ( - f"Global force at position failed.\nExpected forces:\n{expected_forces_mixed}\nGot:\n{composed_force_np}" - ) + global_torque_np = wrench_composer.global_torque_w.numpy() + assert np.allclose(global_torque_np, expected_stored_torques, atol=1e-3, rtol=1e-4) - # Verify torques - composed_torque_np = wrench_composer.composed_torque.numpy() - assert np.allclose(composed_torque_np, expected_torques_mixed, atol=1e-3, rtol=1e-4), ( - f"Global force at position failed.\nExpected torques:\n{expected_torques_mixed}\nGot:\n{composed_torque_np}" - ) + # After compose, output torque should be R^T @ cross(P - link_pos, F) + wrench_composer.compose_to_body_frame() + corrected_torque_w = expected_stored_torques - np.cross(link_pos_np, forces_global_np) + expected_out_torque = quat_rotate_inv_np(link_quat_np, corrected_torque_w) + out_torque_np = wrench_composer.out_torque_b.numpy() + assert np.allclose(out_torque_np, expected_out_torque, atol=1e-3, rtol=1e-4) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) def test_local_vs_global_identity_quaternion(device: str): - """Test that local and global give same result with identity quaternion and zero position.""" + """Test that local and global give same composed result with identity quaternion.""" rng = np.random.default_rng(seed=13) num_envs, num_bodies = 10, 5 - # Create mock with identity pose (default) mock_asset_local = MockRigidObject(num_envs, num_bodies, device) mock_asset_global = MockRigidObject(num_envs, num_bodies, device) wrench_composer_local = WrenchComposer(mock_asset_local) wrench_composer_global = WrenchComposer(mock_asset_global) - # Generate random forces and torques forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) torques_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) torques = wp.from_numpy(torques_np, dtype=wp.vec3f, device=device) - # Apply as local wrench_composer_local.add_forces_and_torques(forces=forces, torques=torques, is_global=False) - - # Apply as global (should be same with identity quaternion) wrench_composer_global.add_forces_and_torques(forces=forces, torques=torques, is_global=True) - # Results should be identical + # Both should produce same body-frame output with identity quaternion + wrench_composer_local.compose_to_body_frame() + wrench_composer_global.compose_to_body_frame() + assert np.allclose( - wrench_composer_local.composed_force.numpy(), - wrench_composer_global.composed_force.numpy(), + wrench_composer_local.out_force_b.numpy(), + wrench_composer_global.out_force_b.numpy(), atol=1e-6, ) assert np.allclose( - wrench_composer_local.composed_torque.numpy(), - wrench_composer_global.composed_torque.numpy(), + wrench_composer_local.out_torque_b.numpy(), + wrench_composer_global.out_torque_b.numpy(), atol=1e-6, ) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) def test_90_degree_rotation_global_force(device: str): - """Test global force with a known 90-degree rotation for easy verification.""" + """Test global force composed to body frame with a known 90-degree rotation.""" num_envs, num_bodies = 1, 1 - # 90-degree rotation around Z-axis: (w, x, y, z) = (cos(45°), 0, 0, sin(45°)) - # This rotates X -> Y, Y -> -X + # 90-degree rotation around Z-axis: body X points along world Y angle = np.pi / 2 link_quat_np = np.array([[[[np.cos(angle / 2), 0, 0, np.sin(angle / 2)]]]], dtype=np.float32).reshape(1, 1, 4) link_quat_torch = torch.from_numpy(link_quat_np) @@ -599,28 +561,27 @@ def test_90_degree_rotation_global_force(device: str): mock_asset = MockRigidObject(num_envs, num_bodies, device, link_quat=link_quat_torch) wrench_composer = WrenchComposer(mock_asset) - # Apply force in global +X direction + # Apply global force in +X direction force_global = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32) force_wp = wp.from_numpy(force_global, dtype=wp.vec3f, device=device) - wrench_composer.add_forces_and_torques(forces=force_wp, is_global=True) - # In mixed representation, global forces stay unchanged - expected_force_mixed = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32) + wrench_composer.compose_to_body_frame() - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, expected_force_mixed, atol=1e-5), ( - f"90-degree rotation test failed.\nExpected:\n{expected_force_mixed}\nGot:\n{composed_force_np}" + # Global +X rotated to body frame: quat_rotate_inv(90deg_Z, (1,0,0)) = (0,-1,0) + expected_force_b = np.array([[[0.0, -1.0, 0.0]]], dtype=np.float32) + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, expected_force_b, atol=1e-5), ( + f"Expected:\n{expected_force_b}\nGot:\n{out_force_np}" ) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) def test_90_degree_rotation_local_force(device: str): - """Test local force with a known 90-degree rotation for easy verification.""" + """Test local force stays unchanged in body frame after compose.""" num_envs, num_bodies = 1, 1 - # 90-degree rotation around Z-axis: (w, x, y, z) = (cos(45°), 0, 0, sin(45°)) - # This rotates X -> Y, Y -> -X + # 90-degree rotation around Z-axis angle = np.pi / 2 link_quat_np = np.array([[[[np.cos(angle / 2), 0, 0, np.sin(angle / 2)]]]], dtype=np.float32).reshape(1, 1, 4) link_quat_torch = torch.from_numpy(link_quat_np) @@ -631,51 +592,48 @@ def test_90_degree_rotation_local_force(device: str): # Apply force in local +X direction force_local = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32) force_wp = wp.from_numpy(force_local, dtype=wp.vec3f, device=device) - wrench_composer.add_forces_and_torques(forces=force_wp, is_global=False) - # In mixed representation, local forces get rotated to global: local +X becomes global +Y - expected_force_mixed = np.array([[[0.0, 1.0, 0.0]]], dtype=np.float32) + wrench_composer.compose_to_body_frame() - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, expected_force_mixed, atol=1e-5), ( - f"90-degree rotation test failed.\nExpected:\n{expected_force_mixed}\nGot:\n{composed_force_np}" + # Local force stays unchanged in body frame + expected_force_b = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32) + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, expected_force_b, atol=1e-5), ( + f"Expected:\n{expected_force_b}\nGot:\n{out_force_np}" ) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) def test_composition_local_and_global(device: str): - """Test that local and global forces can be composed together correctly.""" + """Test that local and global forces compose correctly in body frame.""" rng = np.random.default_rng(seed=14) num_envs, num_bodies = 5, 3 - # Create random link quaternions link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) link_quat_torch = torch.from_numpy(link_quat_np) mock_asset = MockRigidObject(num_envs, num_bodies, device, link_quat=link_quat_torch) wrench_composer = WrenchComposer(mock_asset) - # Generate random local and global forces forces_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) forces_local = wp.from_numpy(forces_local_np, dtype=wp.vec3f, device=device) forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) - # Add local forces first wrench_composer.add_forces_and_torques(forces=forces_local, is_global=False) - - # Add global forces wrench_composer.add_forces_and_torques(forces=forces_global, is_global=True) - # In mixed repr: local forces get rotated to global, global forces stay as-is - local_forces_in_global = quat_rotate_np(link_quat_np, forces_local_np) - expected_total = local_forces_in_global + forces_global_np + wrench_composer.compose_to_body_frame() - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, expected_total, atol=1e-4, rtol=1e-5), ( - f"Mixed local/global composition failed.\nExpected:\n{expected_total}\nGot:\n{composed_force_np}" + # out_force_b = quat_rotate_inv(q, global_force) + local_force + global_in_body = quat_rotate_inv_np(link_quat_np, forces_global_np) + expected_total = global_in_body + forces_local_np + + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, expected_total, atol=1e-4, rtol=1e-5), ( + f"local/global composition failed.\nExpected:\n{expected_total}\nGot:\n{out_force_np}" ) @@ -683,11 +641,10 @@ def test_composition_local_and_global(device: str): @pytest.mark.parametrize("num_envs", [1, 10, 50]) @pytest.mark.parametrize("num_bodies", [1, 3, 5]) def test_local_forces_and_torques_at_local_position(device: str, num_envs: int, num_bodies: int): - """Test local forces at local positions (offset from link frame).""" + """Test local forces at local positions produce correct body-frame output.""" rng = np.random.default_rng(seed=15) for _ in range(5): - # Create random link poses (shouldn't affect local frame calculations) link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) link_pos_torch = torch.from_numpy(link_pos_np) @@ -696,7 +653,6 @@ def test_local_forces_and_torques_at_local_position(device: str, num_envs: int, mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) wrench_composer = WrenchComposer(mock_asset) - # Generate random local forces and local positions (offsets) forces_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) torques_local_np = rng.uniform(-50.0, 50.0, (num_envs, num_bodies, 3)).astype(np.float32) positions_local_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) @@ -704,32 +660,34 @@ def test_local_forces_and_torques_at_local_position(device: str, num_envs: int, torques_local = wp.from_numpy(torques_local_np, dtype=wp.vec3f, device=device) positions_local = wp.from_numpy(positions_local_np, dtype=wp.vec3f, device=device) - # Apply local forces and torques at local positions wrench_composer.add_forces_and_torques( forces=forces_local, torques=torques_local, positions=positions_local, is_global=False ) - # In mixed repr: local forces get rotated to global - expected_forces = quat_rotate_np(link_quat_np, forces_local_np) - # In mixed repr: torque = cross(pos_mixed, force_mixed) + quat_rotate(torques_local) - positions_mixed = quat_rotate_np(link_quat_np, positions_local_np) - expected_torques = np.cross(positions_mixed, expected_forces) + quat_rotate_np(link_quat_np, torques_local_np) + wrench_composer.compose_to_body_frame() + + # Local forces stay in body frame + expected_forces = forces_local_np + # Local torques: cross(local_pos, local_force) + local_torque + expected_torques = np.cross(positions_local_np, forces_local_np) + torques_local_np - # Verify - composed_force_np = wrench_composer.composed_force.numpy() - composed_torque_np = wrench_composer.composed_torque.numpy() + out_force_np = wrench_composer.out_force_b.numpy() + out_torque_np = wrench_composer.out_torque_b.numpy() - assert np.allclose(composed_force_np, expected_forces, atol=1e-3, rtol=1e-5) - assert np.allclose(composed_torque_np, expected_torques, atol=1e-3, rtol=1e-5) + assert np.allclose(out_force_np, expected_forces, atol=1e-3, rtol=1e-5) + assert np.allclose(out_torque_np, expected_torques, atol=1e-3, rtol=1e-5) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_global_force_at_link_origin_no_torque(device: str): - """Test that a global force applied at the link origin produces no torque.""" +def test_global_force_at_link_origin_no_composed_torque(device: str): + """Test that a global force applied at the link origin produces zero composed torque. + + The stored global_torque_w is cross(link_pos, F) (about world origin), but after + compose the correction -cross(link_pos, F) cancels it out, giving zero net torque. + """ rng = np.random.default_rng(seed=16) num_envs, num_bodies = 5, 3 - # Create random link poses link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) link_pos_torch = torch.from_numpy(link_pos_np) @@ -738,78 +696,928 @@ def test_global_force_at_link_origin_no_torque(device: str): mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) wrench_composer = WrenchComposer(mock_asset) - # Generate random global forces forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) - # Position = link position (so offset is zero) + # Position = link position positions_at_link = wp.from_numpy(link_pos_np, dtype=wp.vec3f, device=device) - - # Apply global forces at link origin wrench_composer.add_forces_and_torques(forces=forces_global, positions=positions_at_link, is_global=True) - # In mixed repr: global forces stay unchanged, torque = 0 (since position offset is zero) + # Global force stored unchanged expected_forces = forces_global_np - expected_torques = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + global_force_np = wrench_composer.global_force_w.numpy() + assert np.allclose(global_force_np, expected_forces, atol=1e-4, rtol=1e-5) - composed_force_np = wrench_composer.composed_force.numpy() - composed_torque_np = wrench_composer.composed_torque.numpy() + # Stored torque is cross(link_pos, F), NOT zero + expected_stored_torque = np.cross(link_pos_np, forces_global_np) + global_torque_np = wrench_composer.global_torque_w.numpy() + assert np.allclose(global_torque_np, expected_stored_torque, atol=1e-3, rtol=1e-4) - assert np.allclose(composed_force_np, expected_forces, atol=1e-4, rtol=1e-5) - assert np.allclose(composed_torque_np, expected_torques, atol=1e-4, rtol=1e-5) + # But composed output torque should be zero (correction cancels stored torque) + wrench_composer.compose_to_body_frame() + out_torque_np = wrench_composer.out_torque_b.numpy() + expected_zero = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(out_torque_np, expected_zero, atol=1e-3, rtol=1e-4) @pytest.mark.parametrize("device", ["cuda:0", "cpu"]) -def test_forces_after_asset_pose_change(device: str): - """Test that stored forces accumulate correctly as asset pose changes over multiple iterations. +def test_compose_with_changing_pose(device: str): + """Test that compose_to_body_frame uses current quaternion, not quaternion at set time. - This verifies that the WrenchComposer correctly stores forces in "mixed" representation - (global frame orientation, local frame position) and properly transforms new forces based - on the current pose, while previously stored forces remain unchanged. + This verifies the key behavior: global forces are stored in world frame and only + rotated to body frame at compose time using the current body orientation. """ rng = np.random.default_rng(seed=17) num_envs, num_bodies = 5, 3 - num_iterations = 10 - # Create mock asset with initial pose + # Initial pose + link_quat_np_1 = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_quat_torch_1 = torch.from_numpy(link_quat_np_1) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_quat=link_quat_torch_1) + wrench_composer = WrenchComposer(mock_asset) + + # Set global force + forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces_global = wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=forces_global, is_global=True) + + # Set local force + forces_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces_local = wp.from_numpy(forces_local_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=forces_local, is_global=False) + + # Compose with initial pose + wrench_composer.compose_to_body_frame() + expected_1 = quat_rotate_inv_np(link_quat_np_1, forces_global_np) + forces_local_np + out_1 = wrench_composer.out_force_b.numpy() + assert np.allclose(out_1, expected_1, atol=1e-3, rtol=1e-5) + + # Change the body orientation + link_quat_np_2 = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + mock_asset.data.body_link_quat_w = torch.from_numpy(link_quat_np_2).to(device=device, dtype=torch.float32) + + # Compose again with new pose — should use new quaternion + wrench_composer.compose_to_body_frame() + expected_2 = quat_rotate_inv_np(link_quat_np_2, forces_global_np) + forces_local_np + out_2 = wrench_composer.out_force_b.numpy() + assert np.allclose(out_2, expected_2, atol=1e-3, rtol=1e-5) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_compose_with_changing_position(device: str): + """Test that compose_to_body_frame dynamically adjusts torque based on current position. + + A global force with explicit position produces different composed torque when the body translates. + """ + num_envs, num_bodies = 1, 1 + + # Identity quaternion for simplicity + link_pos_np_1 = np.array([[[1.0, 0.0, 0.0]]], dtype=np.float32) + link_pos_torch_1 = torch.from_numpy(link_pos_np_1) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch_1) + wrench_composer = WrenchComposer(mock_asset) + + # Global force F = (0, 10, 0) with explicit position at world origin + force_np = np.array([[[0.0, 10.0, 0.0]]], dtype=np.float32) + position_np = np.array([[[0.0, 0.0, 0.0]]], dtype=np.float32) + force_wp = wp.from_numpy(force_np, dtype=wp.vec3f, device=device) + position_wp = wp.from_numpy(position_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=force_wp, positions=position_wp, is_global=True) + + # Position at origin → stored torque is cross((0,0,0), F) = 0 + stored_torque = wrench_composer.global_torque_w.numpy() + assert np.allclose(stored_torque, np.zeros((1, 1, 3), dtype=np.float32), atol=1e-6) + + # Compose with link_pos = (1,0,0): correction = -cross((1,0,0), (0,10,0)) = -(0,0,10) = (0,0,-10) + wrench_composer.compose_to_body_frame() + out_torque_1 = wrench_composer.out_torque_b.numpy() + expected_torque_1 = np.array([[[0.0, 0.0, -10.0]]], dtype=np.float32) + assert np.allclose(out_torque_1, expected_torque_1, atol=1e-4), f"Expected {expected_torque_1}, got {out_torque_1}" + + # Move body to origin: link_pos = (0,0,0) + mock_asset.data.body_link_pos_w = torch.zeros((1, 1, 3), dtype=torch.float32, device=device) + + # Recompose: correction = -cross((0,0,0), F) = 0, so composed torque = 0 + wrench_composer._dirty = True + wrench_composer.compose_to_body_frame() + out_torque_2 = wrench_composer.out_torque_b.numpy() + expected_torque_2 = np.zeros((1, 1, 3), dtype=np.float32) + assert np.allclose(out_torque_2, expected_torque_2, atol=1e-4), f"Expected {expected_torque_2}, got {out_torque_2}" + + +# ============================================================================ +# add_raw_buffers_from Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_add_raw_buffers_from(device: str): + """Test that add_raw_buffers_from merges all 5 buffers correctly.""" + rng = np.random.default_rng(seed=20) + num_envs, num_bodies = 5, 3 + + mock_asset = MockRigidObject(num_envs, num_bodies, device) + composer_a = WrenchComposer(mock_asset) + composer_b = WrenchComposer(mock_asset) + + # Add forces to composer_a (local + global without positions → at CoM) + forces_a_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces_a_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_a.add_forces_and_torques( + forces=wp.from_numpy(forces_a_local_np, dtype=wp.vec3f, device=device), is_global=False + ) + composer_a.add_forces_and_torques( + forces=wp.from_numpy(forces_a_global_np, dtype=wp.vec3f, device=device), is_global=True + ) + + # Add forces to composer_b (local + global without positions → at CoM) + forces_b_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces_b_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_b.add_forces_and_torques( + forces=wp.from_numpy(forces_b_local_np, dtype=wp.vec3f, device=device), is_global=False + ) + composer_b.add_forces_and_torques( + forces=wp.from_numpy(forces_b_global_np, dtype=wp.vec3f, device=device), is_global=True + ) + + # Merge b into a + composer_a.add_raw_buffers_from(composer_b) + + # Global forces without positions go to global_force_at_com_w + assert np.allclose(composer_a.global_force_at_com_w.numpy(), forces_a_global_np + forces_b_global_np, atol=1e-4) + assert np.allclose(composer_a.local_force_b.numpy(), forces_a_local_np + forces_b_local_np, atol=1e-4) + + +# ============================================================================ +# Dirty Flag / Warning Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_compose_randomized_cross_product_identity(device: str): + """Test cross(P,F) - cross(link_pos,F) = cross(P - link_pos, F) with random inputs.""" + rng = np.random.default_rng(seed=30) + num_envs, num_bodies = 5, 3 + + for _ in range(5): + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + positions = wp.from_numpy(positions_np, dtype=wp.vec3f, device=device) + + wrench_composer.add_forces_and_torques(forces=forces, positions=positions, is_global=True) + wrench_composer.compose_to_body_frame() + + # Expected: out_torque_b = quat_rotate_inv(q, cross(P - link_pos, F)) + expected_torque_w = np.cross(positions_np - link_pos_np, forces_np) + expected_torque_b = quat_rotate_inv_np(link_quat_np, expected_torque_w) + + out_torque_np = wrench_composer.out_torque_b.numpy() + assert np.allclose(out_torque_np, expected_torque_b, atol=1e-3, rtol=1e-4), ( + f"Cross product identity failed.\nExpected:\n{expected_torque_b}\nGot:\n{out_torque_np}" + ) + + # Also verify force output + expected_force_b = quat_rotate_inv_np(link_quat_np, forces_np) + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, expected_force_b, atol=1e-3, rtol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_multiple_accumulated_global_forces_at_positions(device: str): + """Test that 2+ global forces at different positions accumulate correctly.""" + rng = np.random.default_rng(seed=31) + num_envs, num_bodies = 5, 3 + + for _ in range(5): + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + wrench_composer = WrenchComposer(mock_asset) + + # First force at position P1 + f1_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + p1_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(f1_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(p1_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Second force at position P2 + f2_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + p2_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(f2_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(p2_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Verify stored buffers + expected_stored_force = f1_np + f2_np + expected_stored_torque = np.cross(p1_np, f1_np) + np.cross(p2_np, f2_np) + assert np.allclose(wrench_composer.global_force_w.numpy(), expected_stored_force, atol=1e-3) + assert np.allclose(wrench_composer.global_torque_w.numpy(), expected_stored_torque, atol=1e-3) + + # Compose and verify output + wrench_composer.compose_to_body_frame() + total_force = f1_np + f2_np + corrected_torque_w = expected_stored_torque - np.cross(link_pos_np, total_force) + expected_torque_b = quat_rotate_inv_np(link_quat_np, corrected_torque_w) + expected_force_b = quat_rotate_inv_np(link_quat_np, total_force) + + assert np.allclose(wrench_composer.out_torque_b.numpy(), expected_torque_b, atol=1e-3, rtol=1e-4) + assert np.allclose(wrench_composer.out_force_b.numpy(), expected_force_b, atol=1e-3, rtol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_overwrites_previous_positional_torque(device: str): + """Test that set_forces_and_torques replaces (not accumulates) stored positional torque.""" + rng = np.random.default_rng(seed=32) + num_envs, num_bodies = 5, 3 + + for _ in range(5): + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + # First set + f1_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + p1_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.set_forces_and_torques( + forces=wp.from_numpy(f1_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(p1_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Verify first set stored correctly + expected_torque_1 = np.cross(p1_np, f1_np) + assert np.allclose(wrench_composer.global_torque_w.numpy(), expected_torque_1, atol=1e-3) + assert np.allclose(wrench_composer.global_force_w.numpy(), f1_np, atol=1e-3) + + # Second set should overwrite, not accumulate + f2_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + p2_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.set_forces_and_torques( + forces=wp.from_numpy(f2_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(p2_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Should be ONLY the second set's values + expected_torque_2 = np.cross(p2_np, f2_np) + assert np.allclose(wrench_composer.global_torque_w.numpy(), expected_torque_2, atol=1e-3), ( + "set_forces_and_torques should overwrite, not accumulate positional torque" + ) + assert np.allclose(wrench_composer.global_force_w.numpy(), f2_np, atol=1e-3) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_add_raw_buffers_from_with_positional_torques(device: str): + """Test that add_raw_buffers_from correctly merges composers with positional torques.""" + rng = np.random.default_rng(seed=33) + num_envs, num_bodies = 5, 3 + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) link_pos_torch = torch.from_numpy(link_pos_np) link_quat_torch = torch.from_numpy(link_quat_np) mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + composer_a = WrenchComposer(mock_asset) + composer_b = WrenchComposer(mock_asset) + + # Composer A: global force at position + fa_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + pa_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_a.add_forces_and_torques( + forces=wp.from_numpy(fa_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(pa_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Composer B: global force at position + fb_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + pb_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_b.add_forces_and_torques( + forces=wp.from_numpy(fb_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(pb_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Merge B into A + composer_a.add_raw_buffers_from(composer_b) + + # Verify merged raw buffers + expected_force = fa_np + fb_np + expected_torque = np.cross(pa_np, fa_np) + np.cross(pb_np, fb_np) + assert np.allclose(composer_a.global_force_w.numpy(), expected_force, atol=1e-3) + assert np.allclose(composer_a.global_torque_w.numpy(), expected_torque, atol=1e-3) + + # Compose and verify the output uses combined force for correction + composer_a.compose_to_body_frame() + total_force = fa_np + fb_np + corrected_torque_w = expected_torque - np.cross(link_pos_np, total_force) + expected_torque_b = quat_rotate_inv_np(link_quat_np, corrected_torque_w) + expected_force_b = quat_rotate_inv_np(link_quat_np, total_force) + + assert np.allclose(composer_a.out_torque_b.numpy(), expected_torque_b, atol=1e-3, rtol=1e-4) + assert np.allclose(composer_a.out_force_b.numpy(), expected_force_b, atol=1e-3, rtol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_torque_only_no_correction(device: str): + """Test that global torque without forces gets no correction (cross(link_pos, 0) = 0).""" + rng = np.random.default_rng(seed=34) + num_envs, num_bodies = 5, 3 + + for _ in range(5): + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + wrench_composer = WrenchComposer(mock_asset) + + torques_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.add_forces_and_torques( + torques=wp.from_numpy(torques_np, dtype=wp.vec3f, device=device), is_global=True + ) + + wrench_composer.compose_to_body_frame() + + # No forces → correction is -cross(link_pos, 0) = 0 + # Expected: out_torque_b = quat_rotate_inv(q, T) + expected_torque_b = quat_rotate_inv_np(link_quat_np, torques_np) + out_torque_np = wrench_composer.out_torque_b.numpy() + assert np.allclose(out_torque_np, expected_torque_b, atol=1e-3, rtol=1e-4) + + # Forces should be zero + expected_zero = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(wrench_composer.out_force_b.numpy(), expected_zero, atol=1e-6) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_compose_with_changing_position_and_quaternion(device: str): + """Test that compose adapts when both position and quaternion change simultaneously.""" + rng = np.random.default_rng(seed=35) + num_envs, num_bodies = 5, 3 + + # Initial pose + link_pos_np_1 = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np_1 = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch_1 = torch.from_numpy(link_pos_np_1) + link_quat_torch_1 = torch.from_numpy(link_quat_np_1) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch_1, link_quat=link_quat_torch_1) wrench_composer = WrenchComposer(mock_asset) - # Track expected accumulated force in mixed representation - expected_accumulated_force_np = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + # Add global force at position + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) - for iteration in range(num_iterations): - # Randomly choose whether to apply local or global forces - use_global = rng.choice([True, False]) + # Compose with pose 1 + wrench_composer.compose_to_body_frame() + stored_torque = np.cross(positions_np, forces_np) + corrected_1 = stored_torque - np.cross(link_pos_np_1, forces_np) + expected_torque_1 = quat_rotate_inv_np(link_quat_np_1, corrected_1) + expected_force_1 = quat_rotate_inv_np(link_quat_np_1, forces_np) - # Generate random forces - forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) - forces_wp = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + assert np.allclose(wrench_composer.out_torque_b.numpy(), expected_torque_1, atol=1e-3, rtol=1e-4) + assert np.allclose(wrench_composer.out_force_b.numpy(), expected_force_1, atol=1e-3, rtol=1e-4) - # Apply forces - wrench_composer.add_forces_and_torques(forces=forces_wp, is_global=use_global) + # Change BOTH position and quaternion + link_pos_np_2 = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np_2 = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + mock_asset.data.body_link_pos_w = torch.from_numpy(link_pos_np_2).to(device=device, dtype=torch.float32) + mock_asset.data.body_link_quat_w = torch.from_numpy(link_quat_np_2).to(device=device, dtype=torch.float32) - # Update expected accumulated force based on current pose - if use_global: - # Global forces are stored as-is in mixed representation - expected_accumulated_force_np += forces_np - else: - # Local forces are rotated by current link quaternion to mixed representation - forces_in_mixed = quat_rotate_np(link_quat_np, forces_np) - expected_accumulated_force_np += forces_in_mixed + # Compose with pose 2 + wrench_composer._dirty = True + wrench_composer.compose_to_body_frame() + corrected_2 = stored_torque - np.cross(link_pos_np_2, forces_np) + expected_torque_2 = quat_rotate_inv_np(link_quat_np_2, corrected_2) + expected_force_2 = quat_rotate_inv_np(link_quat_np_2, forces_np) + + assert np.allclose(wrench_composer.out_torque_b.numpy(), expected_torque_2, atol=1e-3, rtol=1e-4) + assert np.allclose(wrench_composer.out_force_b.numpy(), expected_force_2, atol=1e-3, rtol=1e-4) + + # Verify they differ (different pose → different output) + assert not np.allclose(expected_torque_1, expected_torque_2, atol=1e-3) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_large_origin_offset_precision(device: str): + """Test that compose correction doesn't lose unacceptable precision at large world offsets. + + The compose kernel computes: cross(P, F) - cross(link_pos, F) = cross(P - link_pos, F). + When link_pos is large (e.g., 2000), both cross products are ~O(20000) but nearly cancel + to ~O(10), risking catastrophic cancellation in float32. + + This test compares the kernel's result against a direct cross(P - link_pos, F) reference + (which has no cancellation) across increasing world offsets. + """ + rng = np.random.default_rng(seed=50) + num_envs, num_bodies = 10, 3 + + offsets = [0.0, 100.0, 1000.0, 2000.0, 5000.0] + + for world_offset in offsets: + # Random small perturbations around the offset + link_pos_np = rng.uniform(-1.0, 1.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_pos_np[..., 0] += world_offset # shift along X - # Verify composed forces match expected accumulation - composed_force_np = wrench_composer.composed_force.numpy() - assert np.allclose(composed_force_np, expected_accumulated_force_np, atol=1e-3, rtol=1e-5) + # Force application point: 1m relative offset from link + relative_offset_np = rng.uniform(-2.0, 2.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = link_pos_np + relative_offset_np - # Change the link position and orientation for the next iteration - link_pos_np = rng.uniform(-20.0, 20.0, (num_envs, num_bodies, 3)).astype(np.float32) + # Random force and quaternion + forces_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) - mock_asset.data.body_link_pos_w = torch.from_numpy(link_pos_np).to(device=device, dtype=torch.float32) - mock_asset.data.body_link_quat_w = torch.from_numpy(link_quat_np).to(device=device, dtype=torch.float32) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + wrench_composer = WrenchComposer(mock_asset) + + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(forces_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + wrench_composer.compose_to_body_frame() + actual_torque_b = wrench_composer.out_torque_b.numpy() + + # Reference: compute cross(P - link_pos, F) directly (no cancellation) + reference_torque_w = np.cross(relative_offset_np, forces_np) + reference_torque_b = quat_rotate_inv_np(link_quat_np, reference_torque_w) + + # With atol=0.1, we accept ~0.01 rad/s error on a 1kg cube — acceptable for robotics + # but flags gross precision issues from catastrophic cancellation. + assert np.allclose(actual_torque_b, reference_torque_b, atol=0.1, rtol=0.0), ( + f"Precision loss at world offset {world_offset}:\n" + f" max absolute error: {np.max(np.abs(actual_torque_b - reference_torque_b)):.6f}\n" + f" mean absolute error: {np.mean(np.abs(actual_torque_b - reference_torque_b)):.6f}" + ) + + +# ============================================================================ +# Dirty Flag / Warning Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_dirty_flag_warns_on_stale_output_access(device: str): + """Test that accessing output properties without calling compose_to_body_frame() emits a warning.""" + num_envs, num_bodies = 2, 1 + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = np.array([[[1.0, 0.0, 0.0]], [[0.0, 1.0, 0.0]]], dtype=np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=forces, is_global=False) + + # Accessing output without compose_to_body_frame() should warn + with warnings_module.catch_warnings(record=True) as caught: + warnings_module.simplefilter("always") + _ = wrench_composer.out_force_b + assert len(caught) == 1 + assert "compose_to_body_frame()" in str(caught[0].message) + assert "raw buffer properties" in str(caught[0].message) + + # Second access should NOT warn (compose was triggered by first access) + with warnings_module.catch_warnings(record=True) as caught: + warnings_module.simplefilter("always") + _ = wrench_composer.out_force_b + assert len(caught) == 0 + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_no_warning_after_explicit_compose(device: str): + """Test that no warning is emitted when compose_to_body_frame() is called before accessing output.""" + num_envs, num_bodies = 2, 1 + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = np.array([[[1.0, 0.0, 0.0]], [[0.0, 1.0, 0.0]]], dtype=np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=forces, is_global=False) + + wrench_composer.compose_to_body_frame() + + # No warning expected + with warnings_module.catch_warnings(record=True) as caught: + warnings_module.simplefilter("always") + _ = wrench_composer.out_force_b + _ = wrench_composer.out_torque_b + _ = wrench_composer.out_force_b_as_torch + _ = wrench_composer.out_torque_b_as_torch + assert len(caught) == 0 + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_dirty_flag_reset_after_reset(device: str): + """Test that reset clears the dirty flag so no warning is emitted on output access.""" + num_envs, num_bodies = 2, 1 + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = np.array([[[1.0, 0.0, 0.0]], [[0.0, 1.0, 0.0]]], dtype=np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=forces, is_global=False) + + wrench_composer.reset() + + # No warning expected after reset (dirty=False, output is zeroed) + with warnings_module.catch_warnings(record=True) as caught: + warnings_module.simplefilter("always") + out = wrench_composer.out_force_b.numpy() + assert len(caught) == 0 + # Output should be zero after reset + assert np.allclose(out, np.zeros((num_envs, num_bodies, 3), dtype=np.float32)) + + +# ============================================================================ +# Global Force at CoM (No Position) Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_no_position_zero_torque_at_offset(device: str): + """Test that a global force without positions produces zero torque even at large body offsets. + + This is the key regression test for the bug where global forces without positions + would produce spurious torque proportional to the body's distance from the world origin. + """ + rng = np.random.default_rng(seed=40) + num_envs, num_bodies = 5, 3 + + for world_offset in [0.0, 100.0, 1000.0, 2000.0]: + link_pos_np = rng.uniform(-1.0, 1.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_pos_np[..., 0] += world_offset + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=forces, is_global=True) + + wrench_composer.compose_to_body_frame() + + # Force at CoM → no positional torque + out_torque_np = wrench_composer.out_torque_b.numpy() + expected_zero = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(out_torque_np, expected_zero, atol=1e-4), ( + f"Spurious torque at world offset {world_offset}:\n" + f" max torque magnitude: {np.max(np.abs(out_torque_np)):.6f}\n" + f" Expected zero torque for global force without positions." + ) + + # Force should be correctly rotated to body frame + expected_force_b = quat_rotate_inv_np(link_quat_np, forces_np) + out_force_np = wrench_composer.out_force_b.numpy() + assert np.allclose(out_force_np, expected_force_b, atol=1e-3, rtol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_at_com_mixed_with_positional(device: str): + """Test that global forces at CoM and global forces with positions compose correctly together.""" + rng = np.random.default_rng(seed=41) + num_envs, num_bodies = 5, 3 + + for _ in range(5): + link_pos_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_pos_torch = torch.from_numpy(link_pos_np) + link_quat_torch = torch.from_numpy(link_quat_np) + + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_pos=link_pos_torch, link_quat=link_quat_torch) + wrench_composer = WrenchComposer(mock_asset) + + # Force 1: global with position (goes to global_force_w) + f1_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + p1_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(f1_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(p1_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Force 2: global without position (goes to global_force_at_com_w) + f2_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(f2_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + wrench_composer.compose_to_body_frame() + + # Total force in world frame: f1 (positional) + f2 (at CoM) + total_force_w = f1_np + f2_np + expected_force_b = quat_rotate_inv_np(link_quat_np, total_force_w) + + # Torque: only f1 participates in correction + stored_torque = np.cross(p1_np, f1_np) + corrected_torque_w = stored_torque - np.cross(link_pos_np, f1_np) + expected_torque_b = quat_rotate_inv_np(link_quat_np, corrected_torque_w) + + out_force_np = wrench_composer.out_force_b.numpy() + out_torque_np = wrench_composer.out_torque_b.numpy() + + assert np.allclose(out_force_np, expected_force_b, atol=1e-3, rtol=1e-4) + assert np.allclose(out_torque_np, expected_torque_b, atol=1e-3, rtol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_add_raw_buffers_from_with_force_at_com(device: str): + """Test that add_raw_buffers_from correctly merges the global_force_at_com_w buffer.""" + rng = np.random.default_rng(seed=42) + num_envs, num_bodies = 5, 3 + + mock_asset = MockRigidObject(num_envs, num_bodies, device) + composer_a = WrenchComposer(mock_asset) + composer_b = WrenchComposer(mock_asset) + + # Composer A: global force without position (at CoM) + fa_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_a.add_forces_and_torques(forces=wp.from_numpy(fa_np, dtype=wp.vec3f, device=device), is_global=True) + + # Composer B: global force without position (at CoM) + fb_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + composer_b.add_forces_and_torques(forces=wp.from_numpy(fb_np, dtype=wp.vec3f, device=device), is_global=True) + + # Merge B into A + composer_a.add_raw_buffers_from(composer_b) + + # Verify merged global_force_at_com_w buffer + assert np.allclose(composer_a.global_force_at_com_w.numpy(), fa_np + fb_np, atol=1e-4) + + +# ============================================================================ +# API Behavior Tests +# ============================================================================ + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_active_property_lifecycle(device: str): + """Test active property transitions: False → True after add → False after reset.""" + num_envs, num_bodies = 4, 2 + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + # Initially inactive + assert wrench_composer.active is False + + # Active after adding forces + forces_np = np.ones((num_envs, num_bodies, 3), dtype=np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + wrench_composer.add_forces_and_torques(forces=forces) + assert wrench_composer.active is True + + # Inactive after full reset + wrench_composer.reset() + assert wrench_composer.active is False + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_clears_all_buffers(device: str): + """Test that set clears all 5 input buffers before writing new values. + + 1. set(forces=F1, positions=P1, is_global=True) → writes global_force_w + global_torque_w + 2. set(forces=F2, is_global=True) (no positions) → writes global_force_at_com_w, clears global_force_w + 3. Verify global_force_w is zero (cleared), global_force_at_com_w has F2 + 4. Compose and verify output includes only F2 + """ + rng = np.random.default_rng(seed=100) + num_envs, num_bodies = 4, 2 + + link_quat_np = random_unit_quaternion_np(rng, (num_envs, num_bodies)) + link_quat_torch = torch.from_numpy(link_quat_np) + mock_asset = MockRigidObject(num_envs, num_bodies, device, link_quat=link_quat_torch) + wrench_composer = WrenchComposer(mock_asset) + zeros = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + + # Step 1: set with positions → global_force_w + global_torque_w + f1_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + p1_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.set_forces_and_torques( + forces=wp.from_numpy(f1_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(p1_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + assert np.allclose(wrench_composer.global_force_w.numpy(), f1_np, atol=1e-4) + + # Step 2: set without positions → global_force_at_com_w; should clear global_force_w + f2_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + wrench_composer.set_forces_and_torques( + forces=wp.from_numpy(f2_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Step 3: global_force_w cleared, only global_force_at_com_w has data + assert np.allclose(wrench_composer.global_force_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_torque_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_force_at_com_w.numpy(), f2_np, atol=1e-4) + assert np.allclose(wrench_composer.local_force_b.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.local_torque_b.numpy(), zeros, atol=1e-7) + + # Step 4: compose — only F2 contributes (no stale F1) + wrench_composer.compose_to_body_frame() + expected_force_b = quat_rotate_inv_np(link_quat_np, f2_np) + assert np.allclose(wrench_composer.out_force_b.numpy(), expected_force_b, atol=1e-3, rtol=1e-4) + # No positional torque — force at CoM + assert np.allclose(wrench_composer.out_torque_b.numpy(), zeros, atol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_partial_reset_preserves_other_envs(device: str): + """Test that reset(env_ids=[0, 2]) zeros envs 0 and 2 but preserves env 1.""" + rng = np.random.default_rng(seed=101) + num_envs, num_bodies = 3, 2 + + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + # Add forces to all 3 envs (local + global + global at CoM) + forces_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces_global_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces_at_com_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + torques_local_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(forces_local_np, dtype=wp.vec3f, device=device), + torques=wp.from_numpy(torques_local_np, dtype=wp.vec3f, device=device), + is_global=False, + ) + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(forces_global_np, dtype=wp.vec3f, device=device), + positions=wp.from_numpy(positions_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + wrench_composer.add_forces_and_torques( + forces=wp.from_numpy(forces_at_com_np, dtype=wp.vec3f, device=device), + is_global=True, + ) + + # Reset envs 0 and 2 only + reset_ids = wp.array([0, 2], dtype=wp.int32, device=device) + wrench_composer.reset(env_ids=reset_ids) + + zeros = np.zeros((num_bodies, 3), dtype=np.float32) + + # All 7 buffers: envs 0 and 2 should be zeroed + for buf in [ + wrench_composer.global_force_w, + wrench_composer.global_torque_w, + wrench_composer.global_force_at_com_w, + wrench_composer.local_force_b, + wrench_composer.local_torque_b, + wrench_composer._out_force_b, + wrench_composer._out_torque_b, + ]: + buf_np = buf.numpy() + assert np.allclose(buf_np[0], zeros, atol=1e-7), "Env 0 not zeroed in buffer" + assert np.allclose(buf_np[2], zeros, atol=1e-7), "Env 2 not zeroed in buffer" + + # Env 1 should retain its data + assert np.allclose(wrench_composer.local_force_b.numpy()[1], forces_local_np[1], atol=1e-4) + assert np.allclose(wrench_composer.local_torque_b.numpy()[1], torques_local_np[1], atol=1e-4) + assert np.allclose(wrench_composer.global_force_w.numpy()[1], forces_global_np[1], atol=1e-4) + expected_global_torque_1 = np.cross(positions_np[1], forces_global_np[1]) + assert np.allclose(wrench_composer.global_torque_w.numpy()[1], expected_global_torque_1, atol=1e-3) + assert np.allclose(wrench_composer.global_force_at_com_w.numpy()[1], forces_at_com_np[1], atol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_torch_tensor_input_conversion(device: str): + """Test that add_forces_and_torques correctly handles torch.Tensor inputs.""" + rng = np.random.default_rng(seed=102) + num_envs, num_bodies = 4, 2 + + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + torques_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + positions_np = rng.uniform(-10.0, 10.0, (num_envs, num_bodies, 3)).astype(np.float32) + + # Pass torch tensors instead of warp arrays + forces_torch = torch.from_numpy(forces_np).to(device=device) + torques_torch = torch.from_numpy(torques_np).to(device=device) + positions_torch = torch.from_numpy(positions_np).to(device=device) + + wrench_composer.add_forces_and_torques( + forces=forces_torch, torques=torques_torch, positions=positions_torch, is_global=False + ) + + # Verify buffers contain correct values + assert np.allclose(wrench_composer.local_force_b.numpy(), forces_np, atol=1e-4) + expected_torque = torques_np + np.cross(positions_np, forces_np) + assert np.allclose(wrench_composer.local_torque_b.numpy(), expected_torque, atol=1e-3) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_resolve_indices_slice_none(device: str): + """Test that slice(None) selects all envs/bodies.""" + rng = np.random.default_rng(seed=103) + num_envs, num_bodies = 4, 2 + + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + + wrench_composer.add_forces_and_torques(forces=forces, env_ids=slice(None), body_ids=slice(None)) + + assert np.allclose(wrench_composer.local_force_b.numpy(), forces_np, atol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_resolve_indices_invalid_slice_raises(device: str): + """Test that a non-None slice for env_ids raises ValueError.""" + num_envs, num_bodies = 4, 2 + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = np.ones((2, num_bodies, 3), dtype=np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + + with pytest.raises(ValueError, match="Doesn't support slice input"): + wrench_composer.add_forces_and_torques(forces=forces, env_ids=slice(0, 5)) + + with pytest.raises(ValueError, match="Doesn't support slice input"): + wrench_composer.add_forces_and_torques(forces=forces, body_ids=slice(0, 2)) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_set_local_forces(device: str): + """Test that set_forces_and_torques with is_global=False writes to local buffers only.""" + rng = np.random.default_rng(seed=104) + num_envs, num_bodies = 4, 2 + + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + forces_np = rng.uniform(-100.0, 100.0, (num_envs, num_bodies, 3)).astype(np.float32) + forces = wp.from_numpy(forces_np, dtype=wp.vec3f, device=device) + + wrench_composer.set_forces_and_torques(forces=forces, is_global=False) + + # Local buffer should have the forces + assert np.allclose(wrench_composer.local_force_b.numpy(), forces_np, atol=1e-4) + + # Global buffers should remain zero + zeros = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(wrench_composer.global_force_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_torque_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_force_at_com_w.numpy(), zeros, atol=1e-7) + + # Compose and verify output equals local force (identity quat) + wrench_composer.compose_to_body_frame() + assert np.allclose(wrench_composer.out_force_b.numpy(), forces_np, atol=1e-4) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_both_none_is_noop(device: str): + """Test that add_forces_and_torques with forces=None and torques=None is a no-op.""" + num_envs, num_bodies = 4, 2 + mock_asset = MockRigidObject(num_envs, num_bodies, device) + wrench_composer = WrenchComposer(mock_asset) + + wrench_composer.add_forces_and_torques(forces=None, torques=None) + + # Should remain inactive + assert wrench_composer.active is False + + # All buffers should be zero + zeros = np.zeros((num_envs, num_bodies, 3), dtype=np.float32) + assert np.allclose(wrench_composer.local_force_b.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.local_torque_b.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_force_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_torque_w.numpy(), zeros, atol=1e-7) + assert np.allclose(wrench_composer.global_force_at_com_w.numpy(), zeros, atol=1e-7) diff --git a/source/isaaclab/test/utils/test_wrench_composer_integration.py b/source/isaaclab/test/utils/test_wrench_composer_integration.py new file mode 100644 index 00000000000..2831274a4c3 --- /dev/null +++ b/source/isaaclab/test/utils/test_wrench_composer_integration.py @@ -0,0 +1,816 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Integration tests for wrench composer with rigid objects. + +These tests validate that global forces/torques remain invariant under body rotation +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import pytest +import torch + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObject, RigidObjectCfg +from isaaclab.sim import build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + + +def generate_cubes_scene( + num_cubes: int = 1, + height: float = 1.0, + device: str = "cuda:0", +) -> tuple[RigidObject, torch.Tensor]: + """Generate a scene with the provided number of cubes.""" + origins = torch.tensor([(i * 1.0, 0, height) for i in range(num_cubes)]).to(device) + for i, origin in enumerate(origins): + sim_utils.create_prim(f"/World/Table_{i}", "Xform", translation=origin) + + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + ) + + cube_object_cfg = RigidObjectCfg( + prim_path="/World/Table_.*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height)), + ) + cube_object = RigidObject(cfg=cube_object_cfg) + return cube_object, origins + + +N_STEPS = 100 +FORCE_MAGNITUDE = 10.0 +TORQUE_MAGNITUDE = 1.0 + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_invariant_under_rotation(device): + """Test that a permanent global force produces the same acceleration before and after body rotation. + + A global +X force is applied. After 100 steps the body is rotated 180deg about Z. + The acceleration (delta_v per phase) should be the same in both phases because the + force is in the global frame and should not rotate with the body. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + mass = cube_object.root_physx_view.get_masses()[0].item() + com = cube_object.data.body_com_pos_w.clone() + + # Apply permanent global force along +X at CoM + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=com, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run N_STEPS + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase1 = cube_object.data.root_lin_vel_w[0].clone() + + # Rotate body 180deg about Z (quat wxyz = [0, 0, 0, 1]) while keeping velocity + root_pose = cube_object.data.root_state_w[0, :7].clone().unsqueeze(0) + root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # 180deg about Z + cube_object.write_root_pose_to_sim(root_pose) + + # Phase 2: run N_STEPS more + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase2 = cube_object.data.root_lin_vel_w[0].clone() + + # Acceleration should be same in both phases: delta_v_phase2 ≈ delta_v_phase1 + delta_v_phase1 = vel_after_phase1[0].item() # vx after phase 1 + delta_v_phase2 = vel_after_phase2[0].item() - vel_after_phase1[0].item() # vx gained in phase 2 + + expected_dv = FORCE_MAGNITUDE / mass * sim.cfg.dt * N_STEPS + + torch.testing.assert_close( + torch.tensor(delta_v_phase1), + torch.tensor(expected_dv), + rtol=0.001, + atol=0.0001, + ) + torch.testing.assert_close( + torch.tensor(delta_v_phase2), + torch.tensor(expected_dv), + rtol=0.001, + atol=0.0001, + ) + + # Y and Z velocity should remain ~0 + assert abs(vel_after_phase2[1].item()) < 0.5, f"Unexpected Y velocity: {vel_after_phase2[1].item()}" + assert abs(vel_after_phase2[2].item()) < 0.5, f"Unexpected Z velocity: {vel_after_phase2[2].item()}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_local_force_follows_rotation(device): + """Test that a permanent local force rotates with the body. + + A local +X force is applied. After 100 steps the body is rotated 180deg about Z. + Since local +X is now world -X, the force should decelerate the body back towards zero velocity. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Apply permanent local force along body +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=False, + ) + + # Phase 1: run N_STEPS — object accelerates along world +X + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase1 = cube_object.data.root_lin_vel_w[0].clone() + assert vel_after_phase1[0].item() > 1.0, "Object should be moving in +X" + + # Rotate body 180deg about Z while keeping velocity + root_pose = cube_object.data.root_state_w[0, :7].clone().unsqueeze(0) + root_pose[0, 3:7] = torch.tensor([0.0, 0.0, 0.0, 1.0], device=device) # 180deg about Z + cube_object.write_root_pose_to_sim(root_pose) + + # Phase 2: run N_STEPS — local +X is now world -X, so force decelerates + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + vel_after_phase2 = cube_object.data.root_lin_vel_w[0].clone() + + # Velocity should be approximately zero: decelerated by the same amount as it accelerated + torch.testing.assert_close( + vel_after_phase2[0], + torch.tensor(0.0, device=device), + atol=0.0001, + rtol=0.001, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_at_offset_generates_torque(device): + """Test that a global force applied at an offset from CoM generates the expected torque. + + A global +X force applied at +1m Y offset from CoM should produce: + - Linear acceleration in +X + - Angular acceleration about -Z (from cross product: (0,1,0) × (10,0,0) = (0,0,-10)) + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Force at offset: +1m in Y from CoM (global frame) + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE # +X force + + torques = torch.zeros(1, len(body_ids), 3, device=device) + + # Position offset: CoM position + 1m in Y (global frame) + com_pos = cube_object.data.body_com_pos_w[:, body_ids, :3].clone() + positions = com_pos.clone() + positions[..., 1] += 1.0 # +1m Y offset + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + lin_vel = cube_object.data.root_lin_vel_w[0] + ang_vel = cube_object.data.root_ang_vel_w[0] + + # Linear velocity in +X should be positive + assert lin_vel[0].item() > 0.1, f"Expected positive X velocity, got {lin_vel[0].item()}" + + # Angular velocity about Z should be negative (cross product: r × F, r=(0,1,0), F=(10,0,0) -> (0,0,-10)) + assert ang_vel[2].item() < -0.1, f"Expected negative Z angular velocity, got {ang_vel[2].item()}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_torque_invariant_under_rotation(device): + """Test that a permanent global torque produces the same angular acceleration before and after rotation. + + A global +Z torque is applied. After 100 steps the body is rotated 90deg about X. + The angular acceleration (delta_omega per phase) about Z should be the same in both phases + because the torque is in the global frame. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Apply permanent global torque about +Z + forces = torch.zeros(1, len(body_ids), 3, device=device) + torques = torch.zeros(1, len(body_ids), 3, device=device) + torques[..., 2] = TORQUE_MAGNITUDE + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run N_STEPS + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_after_phase1 = cube_object.data.root_ang_vel_w[0, 2].clone().item() + + # Rotate body 90deg about X and zero out velocities so phase 2 starts from rest + # (avoids gyroscopic cross-coupling at high omega) + root_pose = cube_object.data.root_state_w[0, :7].clone().unsqueeze(0) + root_pose[0, 3:7] = torch.tensor([0.7071, 0.7071, 0.0, 0.0], device=device) + cube_object.write_root_pose_to_sim(root_pose) + cube_object.write_root_velocity_to_sim(torch.zeros(1, 6, device=device)) + + # Phase 2: run N_STEPS from rest with different body orientation + for _ in range(N_STEPS): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_after_phase2 = cube_object.data.root_ang_vel_w[0, 2].clone().item() + + # Both phases start from rest — angular acceleration about Z should be the same + torch.testing.assert_close( + torch.tensor(omega_z_after_phase1), + torch.tensor(omega_z_after_phase2), + rtol=0.001, + atol=0.0001, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_torque_after_translation(device): + """Test that global force torque updates dynamically when the body translates. + + Phase 1: Cube at (1,0,0). Global force F=(0,10,0) applied at explicit position (1,0,0). + stored_torque = cross((1,0,0), (0,10,0)) = (0,0,10) + correction = -cross((1,0,0), (0,10,0)) = (0,0,-10) + net torque = 0 → no rotation, only linear acceleration in +Y. + + Phase 2: Teleport cube to origin (0,0,0), zero velocity, don't re-apply force. + stored_torque = (0,0,10) (unchanged in buffer) + correction = -cross((0,0,0), (0,10,0)) = (0,0,0) + net torque = (0,0,10) → rotation about +Z. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Phase 1 setup: Move cube to (1, 0, 1) and apply force at (1, 0, 1) + root_state = cube_object.data.root_state_w.clone() + root_state[0, 0] = 1.0 # x = 1 + root_state[0, 1] = 0.0 # y = 0 + root_state[0, 2] = 1.0 # z = 1 + root_state[0, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) # identity quat + root_state[0, 7:] = 0.0 # zero velocity + cube_object.write_root_state_to_sim(root_state) + + # Step once to let the state settle + sim.step() + cube_object.update(sim.cfg.dt) + + # Get current CoM position for the force application point + com_pos = cube_object.data.body_com_pos_w[:, body_ids, :3].clone() + + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE # +Y force + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=com_pos, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run 50 steps — force at CoM, expect no rotation + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + ang_vel_phase1 = cube_object.data.root_ang_vel_w[0].clone() + lin_vel_phase1 = cube_object.data.root_lin_vel_w[0].clone() + + # Should have linear velocity in +Y + assert lin_vel_phase1[1].item() > 0.1, f"Expected positive Y velocity, got {lin_vel_phase1[1].item()}" + + # Angular velocity should be ~0 (force applied at CoM → no torque) + assert abs(ang_vel_phase1[2].item()) < 0.1, ( + f"Expected ~0 Z angular velocity in phase 1, got {ang_vel_phase1[2].item()}" + ) + + # Phase 2: Teleport cube to origin, zero velocity, don't re-apply force + root_state2 = cube_object.data.root_state_w.clone() + root_state2[0, 0] = 0.0 # x = 0 + root_state2[0, 1] = 0.0 + root_state2[0, 2] = 1.0 # z = 1 + root_state2[0, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state2[0, 7:] = 0.0 # zero velocity + cube_object.write_root_state_to_sim(root_state2) + + # Step once to let state settle + sim.step() + cube_object.update(sim.cfg.dt) + + # Phase 2: run 50 steps — body at origin but stored torque = cross((1,0,1), (0,10,0)) = (-10,0,10) + # correction = -cross((0,0,1), (0,10,0)) = -(0,0,0 - but wait, z=1) + # Actually: stored = cross((com_x,com_y,com_z), (0,10,0)) + # After teleport: correction = -cross(new_pos, F), net torque ≠ 0 since positions differ + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + ang_vel_phase2 = cube_object.data.root_ang_vel_w[0].clone() + + # The X component of position changed from ~1 to ~0, so torque about Z changes. + # stored_torque_z = com_x * Fy = ~1 * 10 = ~10 + # After teleport, correction_z = -new_x * Fy = ~0 * 10 = ~0 + # net torque_z ≈ 10 → positive Z angular velocity + assert ang_vel_phase2[2].item() > 0.5, ( + f"Expected positive Z angular velocity in phase 2, got {ang_vel_phase2[2].item()}" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_torque_reverses_on_opposite_side(device): + """Test that dynamic correction produces correct torque sign depending on body position. + + Phase 1: Cube at (-1, 0, 1). Global F=(0, 10, 0) at world point P=(0, 0, 1). + net torque_z = cross(P - link_pos, F)_z = cross((1,0,0), (0,10,0))_z = +10 + → positive Z angular velocity + + Phase 2: Teleport cube to (+1, 0, 1), zero velocity, don't re-apply force. + net torque_z = cross(P - link_pos, F)_z = cross((-1,0,0), (0,10,0))_z = -10 + → negative Z angular velocity + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Move cube to (-1, 0, 1) + root_state = cube_object.data.root_state_w.clone() + root_state[0, 0] = -1.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state[0, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply permanent global F=(0, 10, 0) at world point P=(0, 0, 1) + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + positions = torch.zeros(1, len(body_ids), 3, device=device) + positions[..., 2] = 1.0 # P = (0, 0, 1) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Phase 1: run 50 steps — expect positive Z angular velocity + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_phase1 = cube_object.data.root_ang_vel_w[0, 2].item() + assert omega_z_phase1 > 0.1, f"Phase 1: expected positive omega_z, got {omega_z_phase1}" + + # Phase 2: Teleport cube to (+1, 0, 1), zero velocity + root_state2 = cube_object.data.root_state_w.clone() + root_state2[0, 0] = 1.0 + root_state2[0, 1] = 0.0 + root_state2[0, 2] = 1.0 + root_state2[0, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state2[0, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state2) + sim.step() + cube_object.update(sim.cfg.dt) + + # Phase 2: run 50 steps — expect negative Z angular velocity + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z_phase2 = cube_object.data.root_ang_vel_w[0, 2].item() + assert omega_z_phase2 < -0.1, f"Phase 2: expected negative omega_z, got {omega_z_phase2}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_no_position_no_torque(device): + """Test that global force without positions produces no torque (applied at CoM). + + A body at (2, 0, 1) with global F=(0, 10, 0) and no positions should experience + only linear acceleration, no rotation. The force is applied at the body's CoM. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Move cube to (2, 0, 1) + root_state = cube_object.data.root_state_w.clone() + root_state[0, 0] = 2.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state[0, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply global F=(0, 10, 0) WITHOUT positions → force at CoM, no torque + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + omega_z = cube_object.data.root_ang_vel_w[0, 2].item() + # No positions → force at CoM → zero torque → zero angular velocity + assert abs(omega_z) < 0.01, f"Expected ~zero omega_z for force at CoM, got {omega_z}" + + # Should still have linear acceleration in +Y + lin_vel_y = cube_object.data.root_lin_vel_w[0, 1].item() + assert lin_vel_y > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel_y}" + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_multi_cube_different_torques_from_same_force(device): + """Test kernel indexing across multiple envs with different CoM positions. + + 2 cubes: Cube 0 at (-1, 0, 1), Cube 1 at (+1, 0, 1). + Same global F=(0, 10, 0) at same world point P=(0, 0, 1) to both cubes. + Cube 0: torque_z = cross((1,0,0), (0,10,0))_z = +10 → omega_z > 0 + Cube 1: torque_z = cross((-1,0,0), (0,10,0))_z = -10 → omega_z < 0 + Both have same linear acceleration in +Y. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=2, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Position cubes: Cube 0 at (-1, 0, 1), Cube 1 at (+1, 0, 1) + root_state = cube_object.data.root_state_w.clone() + root_state[0, 0] = -1.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state[0, 7:] = 0.0 + + root_state[1, 0] = 1.0 + root_state[1, 1] = 0.0 + root_state[1, 2] = 1.0 + root_state[1, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state[1, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply same global F=(0, 10, 0) at P=(0, 0, 1) to both cubes + forces = torch.zeros(2, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE + torques = torch.zeros(2, len(body_ids), 3, device=device) + positions = torch.zeros(2, len(body_ids), 3, device=device) + positions[..., 2] = 1.0 # P = (0, 0, 1) + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Cube 0: omega_z > 0 (force point is to the right of CoM) + omega_z_0 = cube_object.data.root_ang_vel_w[0, 2].item() + assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" + + # Cube 1: omega_z < 0 (force point is to the left of CoM) + omega_z_1 = cube_object.data.root_ang_vel_w[1, 2].item() + assert omega_z_1 < -0.1, f"Cube 1: expected negative omega_z, got {omega_z_1}" + + # Both cubes should have same linear velocity in +Y (same force magnitude) + lin_vel_y_0 = cube_object.data.root_lin_vel_w[0, 1].item() + lin_vel_y_1 = cube_object.data.root_lin_vel_w[1, 1].item() + assert abs(lin_vel_y_0 - lin_vel_y_1) < 0.5, ( + f"Both cubes should have similar Y velocity, got {lin_vel_y_0} and {lin_vel_y_1}" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_global_force_torque_far_from_origin(device): + """Test that global force torque correction produces correct physics at large world coordinates. + + Two cubes with identical relative geometry (force offset = (1, 0, 0) from CoM): + Cube 0 at (0, 0, 1) — near origin (reference) + Cube 1 at (2000, 0, 1) — far from origin + + Both get global F=(0, 10, 0) at offset (1, 0, 0) from their respective CoMs. + Expected torque: cross((1,0,0), (0,10,0)) = (0, 0, 10) for both. + + The compose kernel computes cross(P, F) - cross(link_pos, F): + Cube 0: cross((1,0,1), F) - cross((0,0,1), F) — small values, no cancellation + Cube 1: cross((2001,0,1), F) - cross((2000,0,1), F) — large values nearly cancel + + Both cubes should produce the same angular and linear velocities. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=2, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Position cubes: Cube 0 near origin, Cube 1 far from origin + root_state = cube_object.data.root_state_w.clone() + # Cube 0 at (0, 0, 1) + root_state[0, 0] = 0.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + root_state[0, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state[0, 7:] = 0.0 + # Cube 1 at (2000, 0, 1) + root_state[1, 0] = 2000.0 + root_state[1, 1] = 0.0 + root_state[1, 2] = 1.0 + root_state[1, 3:7] = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device) + root_state[1, 7:] = 0.0 + cube_object.write_root_state_to_sim(root_state) + sim.step() + cube_object.update(sim.cfg.dt) + + # Apply F=(0, 10, 0) at +1m X offset from each cube's CoM + forces = torch.zeros(2, len(body_ids), 3, device=device) + forces[..., 1] = FORCE_MAGNITUDE # +Y force + torques = torch.zeros(2, len(body_ids), 3, device=device) + + # Positions: each cube's CoM + (1, 0, 0) + com_pos = cube_object.data.body_com_pos_w[:, body_ids, :3].clone() + positions = com_pos.clone() + positions[..., 0] += 1.0 # +1m X offset from CoM + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + # Run 50 steps + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Both cubes should have positive omega_z (cross((1,0,0), (0,10,0)) = (0,0,10)) + omega_z_0 = cube_object.data.root_ang_vel_w[0, 2].item() + omega_z_1 = cube_object.data.root_ang_vel_w[1, 2].item() + assert omega_z_0 > 0.1, f"Cube 0: expected positive omega_z, got {omega_z_0}" + assert omega_z_1 > 0.1, f"Cube 1: expected positive omega_z, got {omega_z_1}" + + # omega_z values should match within 1% (same relative geometry) + torch.testing.assert_close( + torch.tensor(omega_z_0), + torch.tensor(omega_z_1), + rtol=0.01, + atol=0.0, + msg=lambda msg: ( + f"Angular velocity mismatch between near-origin and far-from-origin cubes:\n" + f" Cube 0 (near): omega_z = {omega_z_0:.6f}\n" + f" Cube 1 (far): omega_z = {omega_z_1:.6f}\n{msg}" + ), + ) + + # Linear velocity in +Y should also match + lin_vel_y_0 = cube_object.data.root_lin_vel_w[0, 1].item() + lin_vel_y_1 = cube_object.data.root_lin_vel_w[1, 1].item() + torch.testing.assert_close( + torch.tensor(lin_vel_y_0), + torch.tensor(lin_vel_y_1), + rtol=0.01, + atol=0.0, + msg=lambda msg: ( + f"Linear velocity mismatch between near-origin and far-from-origin cubes:\n" + f" Cube 0 (near): lin_vel_y = {lin_vel_y_0:.6f}\n" + f" Cube 1 (far): lin_vel_y = {lin_vel_y_1:.6f}\n{msg}" + ), + ) + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_global_force_no_position_no_rotation_large_offset(device): + """Test that a global force without positions produces no rotation at large offsets. + + A cube is placed at (2000, 0, 1) and a global force F=(0, 10, 0) is applied + without positions. The cube should accelerate linearly but not rotate. + Before the fix, this would produce torque proportional to 2000 and cause rotation. + """ + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True, gravity_enabled=False + ) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Place cube at large X offset + root_state = cube_object.data.default_root_state.clone() + root_state[0, 0] = 2000.0 # large X position + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + cube_object.write_root_pose_to_sim(root_state[:, :7]) + cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + cube_object.reset() + + # Apply global force without positions (should go to CoM, no torque) + forces = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) + forces[0, :, 1] = 10.0 # F_y = 10 N + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + body_ids=body_ids, + is_global=True, + ) + + # Step simulation + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Check: angular velocity should be near zero (no rotation) + ang_vel = cube_object.data.root_ang_vel_w[0] + assert torch.allclose(ang_vel, torch.zeros(3, device=device), atol=0.01), ( + f"Expected near-zero angular velocity, got {ang_vel}. " + "Global force without positions should not produce torque." + ) + + # Check: linear velocity in Y should be positive (force is in +Y) + lin_vel = cube_object.data.root_lin_vel_w[0] + assert lin_vel[1] > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel[1]}" + + +@pytest.mark.parametrize("device", ["cuda:0"]) +def test_global_force_at_com_position_no_rotation_large_offset(device): + """Test that a global force with position at CoM produces no rotation at large offsets. + + A cube is placed at (2000, 0, 1) and a global force F=(0, 10, 0) is applied + at the cube's position (i.e., at its CoM). This should produce zero torque, + serving as a control test alongside test_global_force_no_position_no_rotation_large_offset. + """ + with build_simulation_context( + device=device, add_ground_plane=False, auto_add_lighting=True, gravity_enabled=False + ) as sim: + sim._app_control_on_stop_handle = None + cube_object, _ = generate_cubes_scene(num_cubes=1, height=1.0, device=device) + + sim.reset() + + body_ids, _ = cube_object.find_bodies(".*") + + # Place cube at large X offset + root_state = cube_object.data.default_root_state.clone() + root_state[0, 0] = 2000.0 + root_state[0, 1] = 0.0 + root_state[0, 2] = 1.0 + cube_object.write_root_pose_to_sim(root_state[:, :7]) + cube_object.write_root_velocity_to_sim(root_state[:, 7:]) + cube_object.reset() + + # Apply global force AT the cube's position (torque should cancel) + forces = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) + forces[0, :, 1] = 10.0 + + positions = torch.zeros(cube_object.num_instances, len(body_ids), 3, device=device) + positions[0, :, 0] = 2000.0 + positions[0, :, 2] = 1.0 + + cube_object.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + positions=positions, + body_ids=body_ids, + is_global=True, + ) + + for _ in range(50): + cube_object.write_data_to_sim() + sim.step() + cube_object.update(sim.cfg.dt) + + # Force at CoM → no rotation + ang_vel = cube_object.data.root_ang_vel_w[0] + assert torch.allclose(ang_vel, torch.zeros(3, device=device), atol=0.01), ( + f"Expected near-zero angular velocity, got {ang_vel}. " + "Global force at CoM position should not produce torque." + ) + + lin_vel = cube_object.data.root_lin_vel_w[0] + assert lin_vel[1] > 0.1, f"Expected positive Y velocity from applied force, got {lin_vel[1]}" diff --git a/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py b/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py new file mode 100644 index 00000000000..0f2b3250e45 --- /dev/null +++ b/source/isaaclab/test/utils/test_wrench_composer_vs_physx.py @@ -0,0 +1,682 @@ +# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md). +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Integration tests comparing WrenchComposer output vs raw PhysX apply_forces_and_torques_at_position. + +Two identical rigid objects are placed in the same scene. One uses the WrenchComposer path +(set_forces_and_torques → write_data_to_sim → compose → PhysX apply with is_global=False), +the other uses the raw PhysX API directly (apply_forces_and_torques_at_position with matching +is_global flag). After N steps, both objects should have identical velocities. +""" + +"""Launch Isaac Sim Simulator first.""" + +from isaaclab.app import AppLauncher + +# launch omniverse app +simulation_app = AppLauncher(headless=True).app + +"""Rest everything follows.""" + +import math + +import pytest +import torch + +import isaaclab.sim as sim_utils +from isaaclab.assets import RigidObject, RigidObjectCfg +from isaaclab.sim import build_simulation_context +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR + + +def generate_dual_cube_scene( + num_cubes: int = 1, + height: float = 1.0, + device: str = "cuda:0", + initial_rot: tuple[float, ...] | None = None, + spacing: float = 2.0, +) -> tuple[RigidObject, RigidObject]: + """Generate a scene with two sets of cubes: one for the composer path, one for raw PhysX. + + Both sets share the same spawn config and initial state (except a Y offset to avoid overlap). + + Args: + num_cubes: Number of cubes per group (environments). + height: Spawn height. + device: Simulation device. + initial_rot: Initial quaternion (w, x, y, z). Defaults to identity. + spacing: Distance between env origins in X. Defaults to 2.0. + + Returns: + Tuple of (cube_composer, cube_raw) RigidObject instances. + """ + if initial_rot is None: + initial_rot = (1.0, 0.0, 0.0, 0.0) + + y_offset = max(spacing, 3.0) + + # Create Xform prims for both groups + for i in range(num_cubes): + origin_composer = (i * spacing, 0.0, height) + origin_raw = (i * spacing, y_offset, height) # Y offset to avoid overlap + sim_utils.create_prim(f"/World/Composer_{i}", "Xform", translation=origin_composer) + sim_utils.create_prim(f"/World/Raw_{i}", "Xform", translation=origin_raw) + + spawn_cfg = sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd", + rigid_props=sim_utils.RigidBodyPropertiesCfg(), + ) + + cube_composer_cfg = RigidObjectCfg( + prim_path="/World/Composer_.*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, 0.0, height), rot=initial_rot), + ) + cube_composer = RigidObject(cfg=cube_composer_cfg) + + cube_raw_cfg = RigidObjectCfg( + prim_path="/World/Raw_.*/Object", + spawn=spawn_cfg, + init_state=RigidObjectCfg.InitialStateCfg(pos=(0.0, y_offset, height), rot=initial_rot), + ) + cube_raw = RigidObject(cfg=cube_raw_cfg) + + return cube_composer, cube_raw + + +N_STEPS = 50 +FORCE_MAGNITUDE = 10.0 +TORQUE_MAGNITUDE = 1.0 +# 45 degrees about Z: (cos(22.5°), 0, 0, sin(22.5°)) +ROT_45_Z = (math.cos(math.pi / 8), 0.0, 0.0, math.sin(math.pi / 8)) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_local_force(device): + """Baseline: local force at identity orientation. Composer and raw PhysX should match exactly.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: local force +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=False, + ) + + # Raw PhysX data (flattened for PhysX view API) + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() # no-op (composer inactive) + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=None, + indices=raw_indices, + is_global=False, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Compare velocities + torch.testing.assert_close( + cube_composer.data.root_lin_vel_w, + cube_raw.data.root_lin_vel_w, + rtol=1e-4, + atol=1e-4, + ) + # Both should have ~zero angular velocity (force at CoM, no torque) + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + torch.testing.assert_close( + cube_raw.data.root_ang_vel_w, + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force(device): + """Global force with non-identity rotation (45 deg Z). Rotation matters for frame conversion.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: global force +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Linear velocities should match (same global force, same mass) + torch.testing.assert_close( + cube_composer.data.root_lin_vel_w, + cube_raw.data.root_lin_vel_w, + rtol=1e-4, + atol=1e-4, + ) + # Angular velocities should match + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + cube_raw.data.root_ang_vel_w, + rtol=1e-4, + atol=1e-4, + ) + # Both should have ~zero angular velocity (force at CoM, no torque) + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + torch.testing.assert_close( + cube_raw.data.root_ang_vel_w, + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_local_force_at_position(device): + """Local force at a local offset. Both paths should produce identical cross-product torque.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Local force +X at local offset +0.5m Y + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + positions = torch.zeros(1, len(body_ids), 3, device=device) + positions[..., 1] = 0.5 # +0.5m Y offset in local frame + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=positions, + body_ids=body_ids, + is_global=False, + ) + + # Raw PhysX data (local force at local position) + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_positions = torch.zeros(1, 3, device=device) + raw_positions[:, 1] = 0.5 + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=raw_positions, + indices=raw_indices, + is_global=False, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Both linear and angular velocities should match + torch.testing.assert_close( + cube_composer.data.root_lin_vel_w, + cube_raw.data.root_lin_vel_w, + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + cube_raw.data.root_ang_vel_w, + rtol=1e-4, + atol=1e-4, + ) + + # Sanity: angular velocity should be nonzero (cross-product torque) + assert torch.abs(cube_composer.data.root_ang_vel_w[0, 2]).item() > 0.1, ( + "Expected nonzero Z angular velocity from cross-product torque" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force_at_position(device): + """Global force at world position with non-identity rotation. Both rotation AND position correction matter.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Global force +X + forces = torch.zeros(1, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(1, len(body_ids), 3, device=device) + + # Position = each cube's link_pos + offset (same offset for both) + offset = torch.zeros(1, len(body_ids), 3, device=device) + offset[..., 1] = 1.0 # +1m Y offset in world frame + + pos_composer = cube_composer.data.body_com_pos_w[:, body_ids, :3].clone() + offset + pos_raw = cube_raw.data.body_com_pos_w[:, body_ids, :3].clone() + offset + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + positions=pos_composer, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(1, 3, device=device) + raw_positions = pos_raw.view(-1, 3) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=raw_positions, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Both linear and angular velocities should match + torch.testing.assert_close( + cube_composer.data.root_lin_vel_w, + cube_raw.data.root_lin_vel_w, + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + cube_raw.data.root_ang_vel_w, + rtol=1e-4, + atol=1e-4, + ) + + # Sanity: angular velocity should be nonzero (cross-product torque) + assert torch.abs(cube_composer.data.root_ang_vel_w[0, 2]).item() > 0.1, ( + "Expected nonzero Z angular velocity from positional torque" + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_local_torque(device): + """Local torque at identity orientation. Should produce matching angular velocity.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: local torque about +Z + forces = torch.zeros(1, len(body_ids), 3, device=device) + torques = torch.zeros(1, len(body_ids), 3, device=device) + torques[..., 2] = TORQUE_MAGNITUDE + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=False, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_torques = torch.zeros(1, 3, device=device) + raw_torques[:, 2] = TORQUE_MAGNITUDE + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=None, + indices=raw_indices, + is_global=False, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Angular velocities should match + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + cube_raw.data.root_ang_vel_w, + rtol=1e-4, + atol=1e-4, + ) + # Linear velocity should be ~zero for both (no force) + torch.testing.assert_close( + cube_composer.data.root_lin_vel_w, + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + torch.testing.assert_close( + cube_raw.data.root_lin_vel_w, + torch.zeros(1, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_torque(device): + """Global torque with non-identity rotation (45 deg Z). Composer rotates to body frame internally.""" + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene(num_cubes=1, device=device, initial_rot=ROT_45_Z) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: global torque about +Z + forces = torch.zeros(1, len(body_ids), 3, device=device) + torques = torch.zeros(1, len(body_ids), 3, device=device) + torques[..., 2] = TORQUE_MAGNITUDE + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data + raw_forces = torch.zeros(1, 3, device=device) + raw_torques = torch.zeros(1, 3, device=device) + raw_torques[:, 2] = TORQUE_MAGNITUDE + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Angular velocities should match + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + cube_raw.data.root_ang_vel_w, + rtol=1e-4, + atol=1e-4, + ) + + +NUM_CUBES_MULTI = 4 + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force_multi_env(device): + """Global force (no position) with multiple environments. + + Regression: checks that env-indexing and per-body quaternion handling work correctly + when there is more than one environment. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene( + num_cubes=NUM_CUBES_MULTI, device=device, initial_rot=ROT_45_Z + ) + + sim.reset() + + body_ids, _ = cube_composer.find_bodies(".*") + + # Composer path: global force +X for all envs + forces = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + # Raw PhysX data (one row per env) + raw_forces = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + for _ in range(N_STEPS): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Linear velocities should match across all envs + torch.testing.assert_close( + cube_composer.data.root_lin_vel_w, + cube_raw.data.root_lin_vel_w, + rtol=1e-4, + atol=1e-4, + ) + # Angular velocities should match + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + cube_raw.data.root_ang_vel_w, + rtol=1e-4, + atol=1e-4, + ) + # All envs should have ~zero angular velocity + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + torch.zeros(NUM_CUBES_MULTI, 3, device=device), + rtol=0.0, + atol=1e-4, + ) + + +@pytest.mark.parametrize("device", ["cuda:0", "cpu"]) +def test_composer_vs_physx_global_force_with_reset(device): + """Global force (no position) with a mid-simulation reset of half the envs. + + Regression: after reset the permanent wrench is cleared. Re-setting it should + produce correct behavior even though the object state was just reset. + """ + with build_simulation_context(device=device, gravity_enabled=False, auto_add_lighting=True) as sim: + sim._app_control_on_stop_handle = None + cube_composer, cube_raw = generate_dual_cube_scene( + num_cubes=NUM_CUBES_MULTI, device=device, initial_rot=ROT_45_Z, spacing=20.0 + ) + + sim.reset() + + # Capture initial world-frame state (includes env origin offsets) + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + initial_state_composer = torch.cat( + [ + cube_composer.data.root_link_pos_w, + cube_composer.data.root_link_quat_w, + cube_composer.data.root_com_vel_w, + ], + dim=-1, + ).clone() + initial_state_raw = torch.cat( + [cube_raw.data.root_link_pos_w, cube_raw.data.root_link_quat_w, cube_raw.data.root_com_vel_w], + dim=-1, + ).clone() + + body_ids, _ = cube_composer.find_bodies(".*") + + def apply_global_force(): + """Set the same global +X force on the composer cube.""" + forces = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + forces[..., 0] = FORCE_MAGNITUDE + torques = torch.zeros(NUM_CUBES_MULTI, len(body_ids), 3, device=device) + cube_composer.permanent_wrench_composer.set_forces_and_torques( + forces=forces, + torques=torques, + body_ids=body_ids, + is_global=True, + ) + + apply_global_force() + + # Raw PhysX data + raw_forces = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_forces[:, 0] = FORCE_MAGNITUDE + raw_torques = torch.zeros(NUM_CUBES_MULTI, 3, device=device) + raw_indices = cube_raw._ALL_INDICES + + # Phase 1: run N_STEPS / 2 + half = N_STEPS // 2 + for _ in range(half): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # Reset first half of envs on both cubes + reset_ids = list(range(NUM_CUBES_MULTI // 2)) + reset_ids_torch = torch.tensor(reset_ids, dtype=torch.long, device=device) + + # Reset root state using captured world-frame initial state (includes env origins) + cube_composer.write_root_state_to_sim(initial_state_composer[reset_ids_torch], env_ids=reset_ids_torch) + cube_raw.write_root_state_to_sim(initial_state_raw[reset_ids_torch], env_ids=reset_ids_torch) + + cube_composer.reset(reset_ids) + cube_raw.reset(reset_ids) + + # Re-apply the force (reset cleared the permanent wrench) + apply_global_force() + + # Phase 2: run N_STEPS / 2 more + for _ in range(half): + cube_composer.write_data_to_sim() + cube_raw.write_data_to_sim() + cube_raw.root_physx_view.apply_forces_and_torques_at_position( + force_data=raw_forces, + torque_data=raw_torques, + position_data=None, + indices=raw_indices, + is_global=True, + ) + sim.step() + cube_composer.update(sim.cfg.dt) + cube_raw.update(sim.cfg.dt) + + # All envs: composer vs raw should match + torch.testing.assert_close( + cube_composer.data.root_lin_vel_w, + cube_raw.data.root_lin_vel_w, + rtol=1e-4, + atol=1e-4, + ) + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + cube_raw.data.root_ang_vel_w, + rtol=1e-4, + atol=1e-4, + ) + # All envs should have ~zero angular velocity + torch.testing.assert_close( + cube_composer.data.root_ang_vel_w, + torch.zeros(NUM_CUBES_MULTI, 3, device=device), + rtol=0.0, + atol=1e-4, + )