Fix Newton body-offset Jacobian and streamline core offset math

source/isaaclab/changelog.d/speed-up-task-space-offset-jacobian.changed.rst

@@ -0,0 +1 @@
+* Reduced task-space body-offset Jacobian overhead by using batched cross products instead of skew matrices.

source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py

@@ -265,9 +265,7 @@ def _compute_frame_jacobian(self):
                 self._asset.data.body_quat_w.torch[:, self._body_idx],
             )
             offset_pos_b = math_utils.quat_apply(body_quat_b, self._offset_pos)
-            self._jacobian_b[:, 0:3, :] += torch.bmm(
-                -math_utils.skew_symmetric_matrix(offset_pos_b), self._jacobian_b[:, 3:, :]
-            )
+            self._jacobian_b[:, :3] += torch.linalg.cross(self._jacobian_b[:, 3:], offset_pos_b[:, :, None], dim=1)
 
         return self._jacobian_b
 
@@ -680,9 +678,7 @@ def _compute_ee_jacobian(self):
                 self._asset.data.body_quat_w.torch[:, self._ee_body_idx],
             )
             offset_pos_b = math_utils.quat_apply(body_quat_b, self._offset_pos)
-            self._jacobian_b[:, 0:3, :] += torch.bmm(
-                -math_utils.skew_symmetric_matrix(offset_pos_b), self._jacobian_b[:, 3:, :]
-            )
+            self._jacobian_b[:, :3] += torch.linalg.cross(self._jacobian_b[:, 3:], offset_pos_b[:, :, None], dim=1)
 
     def _compute_ee_pose(self):
         """Computes the pose of the ee frame in root frame."""

source/isaaclab_newton/changelog.d/fix-newton-task-space-body-offset-jacobian.fixed.rst

@@ -0,0 +1,5 @@
+* Fixed the body-offset Jacobian of
+  :class:`~isaaclab_newton.envs.mdp.actions.NewtonDifferentialInverseKinematicsAction` and
+  :class:`~isaaclab_newton.envs.mdp.actions.NewtonOperationalSpaceControllerAction`. The translational rows are now
+  shifted by the offset rotated with the body orientation, and the angular rows are no longer rotated by the offset
+  rotation, matching the target frame's pose and velocity.

source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_task_space_actions.py

@@ -103,15 +103,13 @@ def _compute_ee_jacobian(self) -> torch.Tensor:
         """Compute the target-frame Jacobian in the root frame, shape (num_envs, 6, num_joints)."""
         data = self._asset.data
         jacobian_w = data.body_link_jacobian_w.torch[:, self._jacobi_body_idx, :, self._jacobi_joint_ids]
-        root_rot = math_utils.matrix_from_quat(math_utils.quat_inv(data.root_quat_w.torch))
-        self._jacobian_b[:, :3] = torch.bmm(root_rot, jacobian_w[:, :3])
-        self._jacobian_b[:, 3:] = torch.bmm(root_rot, jacobian_w[:, 3:])
+        jac_lin_w, jac_ang_w = jacobian_w[:, :3], jacobian_w[:, 3:]
         if self._offset_pos is not None:
-            # v_ee = v_link + w_link x r_link_ee, then rotate into the target frame
-            self._jacobian_b[:, :3] += torch.bmm(
-                -math_utils.skew_symmetric_matrix(self._offset_pos), self._jacobian_b[:, 3:]
-            )
-            self._jacobian_b[:, 3:] = torch.bmm(math_utils.matrix_from_quat(self._offset_rot), self._jacobian_b[:, 3:])
+            offset_w = math_utils.quat_apply(data.body_quat_w.torch[:, self._body_idx], self._offset_pos)
+            jac_lin_w = jac_lin_w + torch.linalg.cross(jac_ang_w, offset_w[:, :, None], dim=1)
+        world_to_root = math_utils.matrix_from_quat(math_utils.quat_inv(data.root_quat_w.torch))
+        self._jacobian_b[:, :3] = torch.bmm(world_to_root, jac_lin_w)
+        self._jacobian_b[:, 3:] = torch.bmm(world_to_root, jac_ang_w)
         return self._jacobian_b
 
 

source/isaaclab_newton/test/controllers/test_newton_controllers.py

@@ -3,10 +3,13 @@
 #
 # SPDX-License-Identifier: BSD-3-Clause
 
-"""Tests for the Newton model-free controller wrappers against independent control-law references."""
+"""Newton model-free controllers and the task-space Jacobian supplied to them."""
 
 from __future__ import annotations
 
+import math
+from types import SimpleNamespace
+
 import pytest
 import torch
 from isaaclab_newton.controllers import (
@@ -17,6 +20,9 @@
     NewtonOperationalSpaceController,
     NewtonOperationalSpaceControllerCfg,
 )
+from isaaclab_newton.envs.mdp.actions.newton_task_space_actions import _NewtonTaskSpaceAction
+
+from isaaclab.utils import math as math_utils
 
 _NUM_ENVS = 3
 _NUM_JOINTS = 7
@@ -83,12 +89,38 @@ def test_disabled_input_is_rejected():
 
 
 def test_differential_ik_dls_matches_reference():
-    """Position-only DLS follows ``q + dt * bandwidth * J^T (J J^T + lambda^2 I)^-1 e``."""
+    """An offset-frame Jacobian reaches position-only DLS with the correct frame and control law."""
     torch.manual_seed(0)
     bandwidth, damping, dt = 20.0, 0.1, 0.05
     cfg = NewtonDifferentialIKControllerCfg(bandwidth=bandwidth, damping=damping, axis_weight=(1, 1, 1, 0, 0, 0))
     controller = NewtonDifferentialIKController(cfg, _NUM_ENVS, _NUM_JOINTS, _DEVICE)
-    jacobian, q = _rand(_NUM_ENVS, 6, _NUM_JOINTS), _rand(_NUM_ENVS, _NUM_JOINTS)
+    s = math.sqrt(0.5)
+    root_quat_w = torch.tensor([[s, 0.0, 0.0, s]]).repeat(_NUM_ENVS, 1)
+    body_quat_w = math_utils.quat_mul(root_quat_w, torch.tensor([[0.0, 0.0, s, s]]).repeat(_NUM_ENVS, 1))
+    jacobian_w = _rand(_NUM_ENVS, 1, 6, _NUM_JOINTS)
+    jacobian_w[:, 0, :3, 0] = 0.0
+    jacobian_w[:, 0, 3:, 0] = math_utils.quat_apply(root_quat_w, torch.tensor([[0.0, 0.0, 1.0]]))
+    data = SimpleNamespace(
+        body_link_jacobian_w=SimpleNamespace(torch=jacobian_w),
+        root_quat_w=SimpleNamespace(torch=root_quat_w),
+        body_quat_w=SimpleNamespace(torch=torch.stack([root_quat_w, body_quat_w], dim=1)),
+    )
+    action = SimpleNamespace(
+        _asset=SimpleNamespace(data=data),
+        _body_idx=1,  # fixed-root body 1 uses Jacobian row 0
+        _jacobi_body_idx=0,
+        _jacobi_joint_ids=list(range(_NUM_JOINTS)),
+        _jacobian_b=torch.zeros(_NUM_ENVS, 6, _NUM_JOINTS),
+        _offset_pos=torch.tensor([[1.0, 0.0, 0.0]]).repeat(_NUM_ENVS, 1),
+        _offset_rot=torch.tensor([[s, 0.0, 0.0, s]]).repeat(_NUM_ENVS, 1),
+    )
+    jacobian = _NewtonTaskSpaceAction._compute_ee_jacobian(action)
+    # The body-local x offset is root-local y; a root-z rotation gives linear velocity -x.
+    expected = torch.tensor([[-1.0, 0.0, 0.0, 0.0, 0.0, 1.0]]).repeat(_NUM_ENVS, 1)
+    torch.testing.assert_close(jacobian[:, :, 0], expected, atol=1e-6, rtol=0.0)
+    assert not hasattr(math_utils, "velocity_at_point")
+
+    q = _rand(_NUM_ENVS, _NUM_JOINTS)
     pos, pos_des = _rand(_NUM_ENVS, 3), _rand(_NUM_ENVS, 3)
 
     q_target = controller.compute(_pose(pos), _pose(pos_des), jacobian, q, dt)