Adds body_incoming_joint_wrench_b to ArticulationData field (#2128)

# Description

<!--
Thank you for your interest in sending a pull request. Please make sure
to check the contribution guidelines.

Link:
https://isaac-sim.github.io/IsaacLab/main/source/refs/contributing.html
-->

This PR exposes the physics tensor view of the joint reaction wrenches. 

Fixes #[2127](https://github.com/isaac-sim/IsaacLab/issues/2127)
[1566](https://github.com/isaac-sim/IsaacLab/issues/1566)

<!-- As a practice, it is recommended to open an issue to have
discussions on the proposed pull request.
This makes it easier for the community to keep track of what is being
developed or added, and if a given feature
is demanded by more than one party. -->

## Type of change

<!-- As you go through the list, delete the ones that are not
applicable. -->

- New feature (non-breaking change which adds functionality)

<!--
Example:

| Before | After |
| ------ | ----- |
| _gif/png before_ | _gif/png after_ |

To upload images to a PR -- simply drag and drop an image while in edit
mode and it should upload the image directly. You can then paste that
source into the above before/after sections.
-->

## Checklist

- [ ] I have run the [`pre-commit` checks](https://pre-commit.com/) with
`./isaaclab.sh --format`
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have updated the changelog and the corresponding version in the
extension's `config/extension.toml` file
- [ ] I have added my name to the `CONTRIBUTORS.md` or my name already
exists there

<!--
As you go through the checklist above, you can mark something as done by
putting an x character in it

For example,
- [x] I have done this task
- [ ] I have not done this task
-->

---------
Signed-off-by: James Tigue <166445701+jtigue-bdai@users.noreply.github.com>
Co-authored-by: Mayank Mittal <12863862+Mayankm96@users.noreply.github.com>
Co-authored-by: Kelly Guo <kellyg@nvidia.com>

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

@@ -74,6 +74,7 @@ class ArticulationData:
         self._joint_pos = TimestampedBuffer()
         self._joint_acc = TimestampedBuffer()
         self._joint_vel = TimestampedBuffer()
+        self._body_incoming_joint_wrench_b = TimestampedBuffer()
 
     def update(self, dt: float):
         # update the simulation timestamp
@@ -509,6 +510,22 @@ class ArticulationData:
             self._body_acc_w.timestamp = self._sim_timestamp
         return self._body_acc_w.data
 
+    @property
+    def body_incoming_joint_wrench_b(self) -> torch.Tensor:
+        """Joint reaction wrench applied from body parent to child body in parent body frame.
+
+        Shape is (num_instances, num_bodies, 6). All body reaction wrenches are provided including the root body to the
+        world of an articulation.
+
+        For more information on joint wrenches, please check the`PhysX documentation <https://nvidia-omniverse.github.io/PhysX/physx/5.5.1/docs/Articulations.html#link-incoming-joint-force>`__
+        and the underlying `PhysX Tensor API <https://docs.omniverse.nvidia.com/kit/docs/omni_physics/latest/extensions/runtime/source/omni.physics.tensors/docs/api/python.html#omni.physics.tensors.impl.api.ArticulationView.get_link_incoming_joint_force>`__ .
+        """
+
+        if self._body_incoming_joint_wrench_b.timestamp < self._sim_timestamp:
+            self._body_incoming_joint_wrench_b.data = self._root_physx_view.get_link_incoming_joint_force()
+            self._body_incoming_joint_wrench_b.time_stamp = self._sim_timestamp
+        return self._body_incoming_joint_wrench_b.data
+
     @property
     def projected_gravity_b(self):
         """Projection of the gravity direction on base frame. Shape is (num_instances, 3)."""

source/isaaclab/isaaclab/envs/mdp/observations.py

@@ -233,8 +233,8 @@ def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> tor
     # extract the used quantities (to enable type-hinting)
     asset: Articulation = env.scene[asset_cfg.name]
     # obtain the link incoming forces in world frame
-    link_incoming_forces = asset.root_physx_view.get_link_incoming_joint_force()[:, asset_cfg.body_ids]
-    return link_incoming_forces.view(env.num_envs, -1)
+    body_incoming_joint_wrench_b = asset.data.body_incoming_joint_wrench_b[:, asset_cfg.body_ids]
+    return body_incoming_joint_wrench_b.view(env.num_envs, -1)
 
 
 def imu_orientation(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("imu")) -> torch.Tensor:

source/isaaclab/test/assets/test_articulation.py

@@ -96,10 +96,15 @@ def generate_articulation_cfg(
                     velocity_limit_sim=velocity_limit_sim,
                     effort_limit=effort_limit,
                     velocity_limit=velocity_limit,
-                    stiffness=0.0,
-                    damping=10.0,
+                    stiffness=2000.0,
+                    damping=100.0,
                 ),
             },
+            init_state=ArticulationCfg.InitialStateCfg(
+                pos=(0.0, 0.0, 0.0),
+                joint_pos=({"RevoluteJoint": 1.5708}),
+                rot=(0.7071055, 0.7071081, 0, 0),
+            ),
         )
     elif articulation_type == "single_joint_explicit":
         # we set 80.0 default for max force because default in USD is 10e10 which makes testing annoying.
@@ -1509,6 +1514,96 @@ class TestArticulation(unittest.TestCase):
                                     elif state_location == "link":
                                         torch.testing.assert_close(rand_state, articulation.data.root_link_state_w)
 
+    def test_body_incoming_joint_wrench_b_single_joint(self):
+        """Test the data.body_incoming_joint_wrench_b buffer is populated correctly and statically correct for single joint."""
+        for num_articulations in (2, 1):
+            for device in ("cpu", "cuda:0"):
+                print(num_articulations, device)
+                with self.subTest(num_articulations=num_articulations, device=device):
+                    with build_simulation_context(
+                        gravity_enabled=True, device=device, add_ground_plane=False, auto_add_lighting=True
+                    ) as sim:
+                        sim._app_control_on_stop_handle = None
+                        articulation_cfg = generate_articulation_cfg(articulation_type="single_joint_implicit")
+                        articulation, _ = generate_articulation(
+                            articulation_cfg=articulation_cfg, num_articulations=num_articulations, device=device
+                        )
+
+                        # Play the simulator
+                        sim.reset()
+                        # apply external force
+                        external_force_vector_b = torch.zeros(
+                            (num_articulations, articulation.num_bodies, 3), device=device
+                        )
+                        external_force_vector_b[:, 1, 1] = 10.0  # 10 N in Y direction
+                        external_torque_vector_b = torch.zeros(
+                            (num_articulations, articulation.num_bodies, 3), device=device
+                        )
+                        external_torque_vector_b[:, 1, 2] = 10.0  # 10 Nm in z direction
+
+                        # apply action to the articulation
+                        joint_pos = torch.ones_like(articulation.data.joint_pos) * 1.5708 / 2.0
+                        articulation.write_joint_state_to_sim(
+                            torch.ones_like(articulation.data.joint_pos), torch.zeros_like(articulation.data.joint_vel)
+                        )
+                        articulation.set_joint_position_target(joint_pos)
+                        articulation.write_data_to_sim()
+                        for _ in range(50):
+                            articulation.set_external_force_and_torque(
+                                forces=external_force_vector_b, torques=external_torque_vector_b
+                            )
+                            articulation.write_data_to_sim()
+                            # perform step
+                            sim.step()
+                            # update buffers
+                            articulation.update(sim.cfg.dt)
+
+                            # check shape
+                            self.assertEqual(
+                                articulation.data.body_incoming_joint_wrench_b.shape,
+                                (num_articulations, articulation.num_bodies, 6),
+                            )
+
+                        # calculate expected static
+                        mass = articulation.data.default_mass
+                        pos_w = articulation.data.body_pos_w
+                        quat_w = articulation.data.body_quat_w
+
+                        mass_link2 = mass[:, 1].view(num_articulations, -1)
+                        gravity = (
+                            torch.tensor(sim.cfg.gravity, device="cpu")
+                            .repeat(num_articulations, 1)
+                            .view((num_articulations, 3))
+                        )
+
+                        # NOTE: the com and link pose for single joint are colocated
+                        weight_vector_w = mass_link2 * gravity
+                        # expected wrench from link mass and external wrench
+                        expected_wrench = torch.zeros((num_articulations, 6), device=device)
+                        expected_wrench[:, :3] = math_utils.quat_apply(
+                            math_utils.quat_conjugate(quat_w[:, 0, :]),
+                            weight_vector_w.to(device)
+                            + math_utils.quat_apply(quat_w[:, 1, :], external_force_vector_b[:, 1, :]),
+                        )
+                        expected_wrench[:, 3:] = math_utils.quat_apply(
+                            math_utils.quat_conjugate(quat_w[:, 0, :]),
+                            torch.cross(
+                                pos_w[:, 1, :].to(device) - pos_w[:, 0, :].to(device),
+                                weight_vector_w.to(device)
+                                + math_utils.quat_apply(quat_w[:, 1, :], external_force_vector_b[:, 1, :]),
+                                dim=-1,
+                            )
+                            + math_utils.quat_apply(quat_w[:, 1, :], external_torque_vector_b[:, 1, :]),
+                        )
+
+                        # check value of last joint wrench
+                        torch.testing.assert_close(
+                            expected_wrench,
+                            articulation.data.body_incoming_joint_wrench_b[:, 1, :].squeeze(1),
+                            atol=1e-2,
+                            rtol=1e-3,
+                        )
+
 
 if __name__ == "__main__":
     run_tests()