diff --git a/source/isaaclab/changelog.d/speed-up-task-space-offset-jacobian.changed.rst b/source/isaaclab/changelog.d/speed-up-task-space-offset-jacobian.changed.rst new file mode 100644 index 000000000000..a18330ce10a0 --- /dev/null +++ b/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. diff --git a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py b/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py index afceba8dbd87..038673ab2975 100644 --- a/source/isaaclab/isaaclab/envs/mdp/actions/task_space_actions.py +++ b/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.""" diff --git a/source/isaaclab_newton/changelog.d/fix-newton-task-space-body-offset-jacobian.fixed.rst b/source/isaaclab_newton/changelog.d/fix-newton-task-space-body-offset-jacobian.fixed.rst new file mode 100644 index 000000000000..f3ec05a5269b --- /dev/null +++ b/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. diff --git a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_task_space_actions.py b/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_task_space_actions.py index 026b71bf2fc3..25c24b6c16fd 100644 --- a/source/isaaclab_newton/isaaclab_newton/envs/mdp/actions/newton_task_space_actions.py +++ b/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 diff --git a/source/isaaclab_newton/test/controllers/test_newton_controllers.py b/source/isaaclab_newton/test/controllers/test_newton_controllers.py index 8eb27ea9dc11..61b3fa32af33 100644 --- a/source/isaaclab_newton/test/controllers/test_newton_controllers.py +++ 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)