[BUG FIX] Fix 'RigidJoint.(get_anchor_pos | get_anchor_axis)' getters. (#2012)

Co-authored-by: Alexis DUBURCQ <[email protected]>

Author: alexis779

genesis/engine/entities/rigid_entity/rigid_joint.py

@@ -3,11 +3,11 @@
 
 import genesis as gs
 import genesis.utils.geom as gu
+from genesis.utils import array_class
 from genesis.utils.misc import DeprecationError
 from genesis.repr_base import RBC
 
 
-@ti.data_oriented
 class RigidJoint(RBC):
     """
     Joint class for rigid body entities. Each RigidLink is connected to its parent link via a RigidJoint.
@@ -102,18 +102,11 @@ def get_anchor_pos(self):
         the anchor point is the "output" of the joint transmission, on which the child body is welded.
         """
         tensor = torch.empty((self._solver._B, 3), dtype=gs.tc_float, device=gs.device)
-        self._kernel_get_anchor_pos(tensor)
+        _kernel_get_anchor_pos(self._idx, tensor, self._solver.joints_state)
         if self._solver.n_envs == 0:
             tensor = tensor.squeeze(0)
         return tensor
 
-    @ti.kernel
-    def _kernel_get_anchor_pos(self, tensor: ti.types.ndarray()):
-        for i_b in range(self._solver._B):
-            xpos = self._solver.joints_state.xanchor[self._idx, i_b]
-            for i in ti.static(range(3)):
-                tensor[i_b, i] = xpos[i]
-
     @gs.assert_built
     def get_anchor_axis(self):
         """
@@ -122,18 +115,11 @@ def get_anchor_axis(self):
         See `RigidJoint.get_anchor_pos` documentation for details about the notion on anchor point.
         """
         tensor = torch.empty((self._solver._B, 3), dtype=gs.tc_float, device=gs.device)
-        self._kernel_get_anchor_axis(tensor)
+        _kernel_get_anchor_axis(self._idx, tensor, self._solver.joints_state)
         if self._solver.n_envs == 0:
             tensor = tensor.squeeze(0)
         return tensor
 
-    @ti.kernel
-    def _kernel_get_anchor_axis(self, tensor: ti.types.ndarray()):
-        for i_b in range(self._solver._B):
-            xaxis = self._solver.joints_state.xaxis[self._idx, i_b]
-            for i in ti.static(range(3)):
-                tensor[i_b, i] = xaxis[i]
-
     def set_sol_params(self, sol_params):
         """
         Set the solver parameters of this joint.
@@ -454,3 +440,21 @@ def is_built(self):
 
     def _repr_brief(self):
         return f"{(self._repr_type())}: {self._uid}, name: '{self._name}', idx: {self._idx}, type: {self._type}"
+
+
+@ti.kernel
+def _kernel_get_anchor_pos(joint_idx: ti.i32, tensor: ti.types.ndarray(), joints_state: array_class.JointsState):
+    _B = joints_state.xanchor.shape[1]
+    for i_b in range(_B):
+        xpos = joints_state.xanchor[joint_idx, i_b]
+        for i in ti.static(range(3)):
+            tensor[i_b, i] = xpos[i]
+
+
+@ti.kernel
+def _kernel_get_anchor_axis(joint_idx: ti.i32, tensor: ti.types.ndarray(), joints_state: array_class.JointsState):
+    _B = joints_state.xaxis.shape[1]
+    for i_b in range(_B):
+        xaxis = joints_state.xaxis[joint_idx, i_b]
+        for i in ti.static(range(3)):
+            tensor[i_b, i] = xaxis[i]

tests/test_rigid_physics.py

@@ -3392,3 +3392,30 @@ def test_reset_control(robot_path, tol):
     new_control_force = robot.get_dofs_control_force()
     assert old_control_force.abs().max() > gs.EPS
     assert_allclose(new_control_force, 0, tol=gs.EPS)
+
+
+@pytest.mark.required
+@pytest.mark.parametrize("n_envs", [0, 2])
+def test_joint_get_anchor_pos_and_axis(n_envs):
+    scene = gs.Scene(
+        show_viewer=False,
+        show_FPS=False,
+    )
+    robot = scene.add_entity(
+        gs.morphs.MJCF(
+            file="xml/franka_emika_panda/panda.xml",
+        ),
+    )
+    scene.build(n_envs=n_envs)
+    batch_shape = (n_envs,) if n_envs > 0 else ()
+
+    joint = robot.joints[1]
+    anchor_pos = joint.get_anchor_pos()
+    assert anchor_pos.shape == (*batch_shape, 3)
+    expected_pos = scene.rigid_solver.joints_state.xanchor.to_numpy()
+    assert_allclose(anchor_pos, expected_pos[joint.idx], tol=gs.EPS)
+
+    anchor_axis = joint.get_anchor_axis()
+    assert anchor_axis.shape == (*batch_shape, 3)
+    expected_axis = scene.rigid_solver.joints_state.xaxis.to_numpy()
+    assert_allclose(anchor_axis, expected_axis[joint.idx], tol=gs.EPS)